1. 云端訓(xùn)練任務(wù)為什么會“跑著跑著就沒了”但凡把 PyTorch 訓(xùn)練任務(wù)放到云端 GPU 上跑過的人大概率都經(jīng)歷過這種崩潰時(shí)刻模型訓(xùn)到第 8 個 epochloss 曲線正漂亮地往下走你關(guān)掉瀏覽器去吃了個飯回來一看——SSH 會話斷了進(jìn)程沒了顯存也釋放了日志停在最后一個 step 上前面幾個小時(shí)的算力全打了水漂。更氣人的是你甚至不知道它是被 OOM 殺的、被會話超時(shí)斷的還是被平臺調(diào)度回收的。這個問題的本質(zhì)其實(shí)不是 PyTorch 的鍋而是訓(xùn)練進(jìn)程的生命周期管理沒做好。云端 GPU 任務(wù)面臨三重威脅第一是 SSH 會話本身不持久網(wǎng)絡(luò)抖動、本地休眠、終端關(guān)閉都會觸發(fā) SIGHUP 把前臺進(jìn)程帶走第二是訓(xùn)練腳本自身沒有容錯一旦中途異常退出只能從頭再來第三是平臺側(cè)的資源回收策略比如搶占式實(shí)例、配額凍結(jié)、節(jié)點(diǎn)維護(hù)這些你控制不了但可以提前防御。所以一套能“穩(wěn)定跑完”的方案必須同時(shí)解決兩件事進(jìn)程要活得住后臺?；顮顟B(tài)要存得下斷點(diǎn)續(xù)訓(xùn)。前者靠 tmux 這類會話保持工具后者靠 checkpoint 機(jī)制。這兩件事單獨(dú)看都不復(fù)雜但真正落地到實(shí)際項(xiàng)目里坑非常多——checkpoint 存什么、多久存一次、存到哪、怎么保證原子性、恢復(fù)時(shí)怎么對齊 optimizer 狀態(tài)和 scheduler 步數(shù)每一步都有講究。這篇文章面向的是已經(jīng)在用 PyTorch 做訓(xùn)練、并且需要把任務(wù)放到云端 GPU 上長時(shí)間運(yùn)行的開發(fā)者。不管你是微調(diào)大模型、跑強(qiáng)化學(xué)習(xí)、還是訓(xùn)一個中等規(guī)模的 CV 模型只要單次訓(xùn)練超過半小時(shí)這套東西就值得認(rèn)真搭一遍。下面我會從進(jìn)程保活講到 checkpoint 設(shè)計(jì)再到恢復(fù)邏輯的完整實(shí)現(xiàn)最后聊聊那些只有踩過才知道的細(xì)節(jié)。2. tmux ?；钭層?xùn)練進(jìn)程脫離 SSH 會話獨(dú)立存活2.1 為什么 nohup 不夠用tmux 才是正解很多人第一反應(yīng)是用nohup python train.py 把任務(wù)丟到后臺。這招在簡單場景下能用但一旦你需要中途查看訓(xùn)練狀態(tài)、臨時(shí)敲個命令、或者進(jìn)程卡住時(shí)進(jìn)去調(diào)試nohup 就抓瞎了——你沒法重新“attach”回那個進(jìn)程的終端上下文只能靠日志文件盲猜。tmux 的核心價(jià)值在于它提供了一個持久化的偽終端會話。你在這個會話里啟動的進(jìn)程其父進(jìn)程是 tmux server而不是你的 SSH 連接。SSH 斷了tmux server 還在進(jìn)程就還在。等你重新連上服務(wù)器tmux attach一下屏幕上的輸出原封不動地回來了就像從沒離開過。我實(shí)測下來tmux 相比 screen 更值得推薦原因是它的配置更靈活、窗口分屏更順手、腳本化能力更強(qiáng)。下面是一套我常用的配置和操作流程。2.2 從零搭一個抗斷線的訓(xùn)練會話先裝 tmuxUbuntu/Debian 系直接apt install tmuxCentOS 系yum install tmux。裝完之后建議先寫一份~/.tmux.conf把一些默認(rèn)反人類的快捷鍵改掉# ~/.tmux.conf set -g mouse on # 允許鼠標(biāo)滾輪翻頁、點(diǎn)擊切窗格 set -g history-limit 50000 # 回滾緩沖區(qū)加大方便翻訓(xùn)練日志 set -g base-index 1 # 窗口編號從1開始 setw -g pane-base-index 1 set -g renumber-windows on set -g escape-time 10 # 降低ESC延遲vim用戶友好 bind r source-file ~/.tmux.conf \; display config reloaded配好之后標(biāo)準(zhǔn)操作流程是這樣的# 1. 新建一個名為 train 的會話直接進(jìn)入 tmux new -s train # 2. 在會話里激活環(huán)境并啟動訓(xùn)練 conda activate myenv cd /path/to/project python train.py --config configs/base.yaml 21 | tee logs/train_$(date %m%d_%H%M).log # 3. 按 Ctrlb 然后按 ddetach 出來進(jìn)程繼續(xù)跑 # 4. 下次連上服務(wù)器后重新進(jìn)入 tmux attach -t train這里有個細(xì)節(jié)值得說21 | tee這個組合把 stderr 合并進(jìn) stdout同時(shí)輸出到終端和日志文件。為什么要這么做因?yàn)?PyTorch 的 tqdm 進(jìn)度條、警告信息很多是走 stderr 的不合并的話日志文件里會缺東西。而 tee 保證你既能在 tmux 里實(shí)時(shí)看又有落盤記錄可查。2.3 會話管理里那些容易翻車的點(diǎn)第一個坑是會話名沖突。如果你習(xí)慣用tmux new -s train第二次執(zhí)行會報(bào) duplicate session。正確做法是先tmux ls看看有沒有同名會話有就 attach沒有才 new。我一般寫個小函數(shù)丟進(jìn).bashrctm() { tmux attach -t $1 2/dev/null || tmux new -s $1 }這樣tm train一條命令搞定“有則進(jìn)、無則建”。第二個坑是tmux 里的進(jìn)程被 OOM Killer 干掉。tmux 只能防會話斷開防不了系統(tǒng)內(nèi)存不足時(shí)內(nèi)核殺進(jìn)程。這種情況日志里通?？床坏?Python 的 traceback進(jìn)程直接消失。判斷方法dmesg | grep -i killed process如果看到你的 python 進(jìn)程被 kill那就是系統(tǒng)內(nèi)存問題得從 batch size、數(shù)據(jù)加載器 worker 數(shù)量上找原因。第三個坑是服務(wù)器重啟。tmux 會話不會在重啟后自動恢復(fù)這是物理限制。如果你的平臺會不定期維護(hù)重啟那必須配合 checkpoint 才能扛過去——這也正好引出下一部分。提示tmux 會話是綁定到具體機(jī)器的。如果你用的是 K8s 或容器化平臺Pod 重建后 tmux 會話同樣會丟這時(shí)候保活要靠平臺自身的重啟策略而狀態(tài)恢復(fù)只能靠 checkpoint。3. Checkpoint 到底該存什么不只是 model.state_dict()3.1 一個“能真正恢復(fù)”的 checkpoint 包含哪些字段新手最容易犯的錯就是只存model.state_dict()。結(jié)果恢復(fù)訓(xùn)練時(shí)發(fā)現(xiàn)optimizer 的動量沒了、學(xué)習(xí)率調(diào)度器的步數(shù)歸零了、epoch 計(jì)數(shù)從頭開始了。這樣恢復(fù)出來的訓(xùn)練loss 會有一個明顯的跳變甚至可能發(fā)散。一個完整的、可無縫續(xù)訓(xùn)的 checkpoint至少應(yīng)該包含以下內(nèi)容字段作用不存的后果model_state_dict模型權(quán)重等于從頭訓(xùn)optimizer_state_dict優(yōu)化器動量/二階矩loss 跳變收斂變慢scheduler_state_dict學(xué)習(xí)率調(diào)度狀態(tài)學(xué)習(xí)率重置訓(xùn)練節(jié)奏亂epoch/global_step進(jìn)度計(jì)數(shù)日志、斷點(diǎn)判斷錯亂best_metric歷史最優(yōu)指標(biāo)無法正確保存最優(yōu)模型rng_state可選隨機(jī)數(shù)狀態(tài)數(shù)據(jù)增強(qiáng)、dropout 不可復(fù)現(xiàn)scaler_state_dictAMP混合精度縮放因子AMP 訓(xùn)練恢復(fù)后可能溢出我一般會封裝成一個函數(shù)把該存的都存進(jìn)去import torch import random import numpy as np def save_checkpoint(path, model, optimizer, scheduler, epoch, global_step, best_metric, scalerNone): ckpt { model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict() if scheduler else None, epoch: epoch, global_step: global_step, best_metric: best_metric, rng_state: { torch: torch.get_rng_state(), cuda: torch.cuda.get_rng_state_all(), numpy: np.random.get_state(), python: random.getstate(), }, } if scaler is not None: ckpt[scaler] scaler.state_dict() # 先寫臨時(shí)文件再原子重命名防止寫一半崩潰導(dǎo)致文件損壞 tmp_path path .tmp torch.save(ckpt, tmp_path) os.replace(tmp_path, path)注意最后那個os.replace的寫法。這是原子寫的關(guān)鍵如果直接torch.save(ckpt, path)在寫入過程中進(jìn)程被 kill你會得到一個半截的、無法加載的損壞文件而它恰好覆蓋了上一個好的 checkpoint。用臨時(shí)文件 原子重命名就能保證磁盤上永遠(yuǎn)是一個完整可用的版本。3.2 存盤頻率太勤傷 IO太懶丟進(jìn)度checkpoint 的保存間隔是個權(quán)衡。存太頻繁每個 epoch 都寫幾個 GBIO 成為瓶頸訓(xùn)練速度明顯下降存太稀疏比如 10 個 epoch 才存一次一旦崩潰就丟 10 個 epoch 的進(jìn)度。我的經(jīng)驗(yàn)是分兩檔周期性 checkpoint和最優(yōu) checkpoint。周期性的一般按時(shí)間或步數(shù)觸發(fā)比如每 30 分鐘或每 1000 個 step 存一次覆蓋寫同一個文件或保留最近 2-3 個滾動版本最優(yōu) checkpoint 則在驗(yàn)證指標(biāo)刷新時(shí)保存單獨(dú)命名永不覆蓋。SAVE_INTERVAL_SEC 1800 # 30分鐘 last_save_time time.time() for epoch in range(start_epoch, num_epochs): for step, batch in enumerate(loader): train_step(batch) global_step 1 if time.time() - last_save_time SAVE_INTERVAL_SEC: save_checkpoint(ckpt/last.pt, model, optimizer, scheduler, epoch, global_step, best_metric, scaler) last_save_time time.time()按時(shí)間觸發(fā)比按 step 觸發(fā)更穩(wěn)因?yàn)椴煌?batch 的處理耗時(shí)可能差異很大尤其是變長序列任務(wù)按 step 算不準(zhǔn)真實(shí)的時(shí)間成本。3.3 大模型場景下的存儲優(yōu)化如果你在微調(diào)大模型一個 checkpoint 動輒幾十 GB存盤本身就是個負(fù)擔(dān)。幾個實(shí)用的優(yōu)化方向只存可訓(xùn)練參數(shù)LoRA 微調(diào)時(shí)凍結(jié)的主干權(quán)重根本不用存只存 adapter 的幾十 MB 即可。分片保存用torch.save配合state_dict分片或者直接用safetensors格式加載更快也更安全。異步寫盤把torch.save丟到一個后臺線程里訓(xùn)練主循環(huán)不阻塞。但要注意異步寫盤期間如果進(jìn)程崩潰可能丟最后一次 checkpoint所以關(guān)鍵節(jié)點(diǎn)還是同步寫。import threading def async_save(path, ckpt): def _save(): tmp path .tmp torch.save(ckpt, tmp) os.replace(tmp, path) t threading.Thread(target_save) t.start() return t注意異步寫盤時(shí)ckpt里的 tensor 如果在寫盤期間被訓(xùn)練循環(huán)修改會存出錯亂的數(shù)據(jù)。穩(wěn)妥做法是先copy.deepcopy或者把 state_dict 轉(zhuǎn)到 CPU 再交給后臺線程。4. 恢復(fù)邏輯讓訓(xùn)練從斷點(diǎn)“無縫”接上4.1 加載順序錯了恢復(fù)就白搭恢復(fù)訓(xùn)練時(shí)加載順序是有講究的。正確的順序是先建模型和優(yōu)化器實(shí)例再加載 state_dict最后把 RNG 狀態(tài)也恢復(fù)上。很多人把 RNG 恢復(fù)漏了導(dǎo)致數(shù)據(jù)增強(qiáng)的隨機(jī)序列和中斷前對不上雖然不影響收斂但嚴(yán)格來說不算“無縫”。def load_checkpoint(path, model, optimizer, scheduler, scalerNone): ckpt torch.load(path, map_locationcpu) model.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer]) if scheduler and ckpt.get(scheduler): scheduler.load_state_dict(ckpt[scheduler]) if scaler and ckpt.get(scaler): scaler.load_state_dict(ckpt[scaler]) # 恢復(fù)隨機(jī)數(shù)狀態(tài) rng ckpt.get(rng_state) if rng: torch.set_rng_state(rng[torch]) torch.cuda.set_rng_state_all(rng[cuda]) np.random.set_state(rng[numpy]) random.setstate(rng[python]) return ckpt[epoch], ckpt[global_step], ckpt[best_metric]這里有個非常隱蔽的坑torch.load默認(rèn)會把 tensor 加載到保存時(shí)的設(shè)備上。如果你在 GPU 上存的 checkpoint換到另一臺機(jī)器或 CPU 上加載會直接報(bào)錯。所以務(wù)必加map_locationcpu加載完再model.to(device)。這個參數(shù)看起來不起眼但在跨環(huán)境恢復(fù)時(shí)能救命。4.2 數(shù)據(jù)加載器的“斷點(diǎn)對齊”問題模型狀態(tài)恢復(fù)了但數(shù)據(jù)加載器不一定能對齊。如果你用的是DataLoader的隨機(jī) sampler恢復(fù)后它會從頭開始 shuffle導(dǎo)致某些樣本被重復(fù)訓(xùn)練、某些被跳過。對于 epoch 級別的恢復(fù)這個問題影響不大但如果你做的是 step 級別的精確恢復(fù)就需要保存 sampler 的狀態(tài)。PyTorch 的RandomSampler支持通過set_epoch配合DistributedSampler來保證每個 epoch 的 shuffle 可復(fù)現(xiàn)。單機(jī)場景下更簡單的做法是記錄global_step恢復(fù)時(shí)用itertools.islice跳過已經(jīng)訓(xùn)練過的 batchfrom itertools import islice for epoch in range(start_epoch, num_epochs): loader_iter iter(loader) if epoch start_epoch and global_step 0: # 跳過本 epoch 內(nèi)已經(jīng)訓(xùn)練過的 step steps_done_in_epoch global_step % len(loader) loader_iter islice(loader_iter, steps_done_in_epoch, None) for batch in loader_iter: ...這個邏輯在單機(jī)訓(xùn)練里夠用分布式場景下要復(fù)雜一些需要每個 rank 各自對齊。4.3 自動檢測并恢復(fù)讓腳本自己決定從哪開始最省心的方案是讓訓(xùn)練腳本啟動時(shí)自動檢查有沒有可用的 checkpoint有就從最新的恢復(fù)沒有就從頭開始。這樣配合 tmux即使進(jìn)程意外退出你只要重新跑一遍同樣的命令它自己就接上了。def find_latest_checkpoint(ckpt_dir): if not os.path.isdir(ckpt_dir): return None files [f for f in os.listdir(ckpt_dir) if f.endswith(.pt)] if not files: return None files.sort(keylambda f: os.path.getmtime(os.path.join(ckpt_dir, f))) return os.path.join(ckpt_dir, files[-1]) # 主流程 ckpt_path find_latest_checkpoint(ckpt) if ckpt_path: start_epoch, global_step, best_metric load_checkpoint( ckpt_path, model, optimizer, scheduler, scaler) print(fResumed from {ckpt_path} at epoch {start_epoch}, step {global_step}) else: start_epoch, global_step, best_metric 0, 0, float(inf)配合一個while循環(huán)包住訓(xùn)練主體還能實(shí)現(xiàn)“崩潰自動重啟”while true; do python train.py --config configs/base.yaml code$? if [ $code -eq 0 ]; then echo Training finished normally. break fi echo Training crashed with code $code, restarting in 10s... sleep 10 done這段 shell 邏輯放在 tmux 會話里跑就形成了一個相當(dāng)健壯的“自愈”訓(xùn)練環(huán)境進(jìn)程崩了自動重啟重啟后自動從最新 checkpoint 恢復(fù)tmux 保證會話不斷。我實(shí)測下來這套組合能扛住絕大多數(shù)非硬件級別的意外。5. 那些只有踩過才知道的細(xì)節(jié)5.1 checkpoint 文件損壞與版本兼容前面提到的原子寫能防大部分損壞但還有一種情況磁盤寫滿。當(dāng)磁盤空間不足時(shí)torch.save可能寫出一個截?cái)嗟奈募鴒s.replace依然會成功于是好的 checkpoint 被壞文件覆蓋了。防御方法是保存前檢查磁盤剩余空間import shutil def safe_save(path, ckpt, min_free_gb5): free_gb shutil.disk_usage(os.path.dirname(path)).free / 1e9 if free_gb min_free_gb: print(fWARNING: only {free_gb:.1f}GB free, skip saving {path}) return False tmp path .tmp torch.save(ckpt, tmp) os.replace(tmp, path) return True另一個坑是PyTorch 版本兼容。用 2.0 存的 checkpoint在 1.13 上加載可能因?yàn)閣eights_only默認(rèn)值變化而報(bào)錯?？绨姹净謴?fù)時(shí)建議顯式指定torch.load(path, map_locationcpu, weights_onlyFalse)并確保兩邊的模型定義代碼一致。5.2 GPU 顯存碎片與長時(shí)間運(yùn)行的穩(wěn)定性訓(xùn)練跑得越久越容易遇到顯存碎片問題。表現(xiàn)是明明還有幾 GB 空閑顯存但分配一個小 tensor 就 OOM。這通常是因?yàn)轭l繁的變長分配導(dǎo)致顯存池碎片化。緩解手段有兩個一是設(shè)置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True讓分配器用可擴(kuò)展段減少碎片二是定期在驗(yàn)證階段torch.cuda.empty_cache()但注意這個操作會拖慢速度別在訓(xùn)練循環(huán)里頻繁調(diào)用。還有一個容易被忽略的點(diǎn)長時(shí)間訓(xùn)練后 GPU 溫度墻導(dǎo)致的降頻。如果你發(fā)現(xiàn)訓(xùn)練速度在幾個小時(shí)后明顯變慢但 loss 正常很可能是 GPU 過熱降頻了。用nvidia-smi -q -d TEMPERATURE看看溫度必要時(shí)改善散熱或降低功耗上限。5.3 日志與監(jiān)控出問題時(shí)能快速定位訓(xùn)練崩了不可怕可怕的是不知道為啥崩。我的習(xí)慣是至少記錄三類信息訓(xùn)練指標(biāo)loss、lr、grad_norm、系統(tǒng)指標(biāo)GPU 利用率、顯存、溫度、異常堆棧完整 traceback 落盤。grad_norm 尤其值得記錄它是判斷訓(xùn)練是否健康的早期信號。如果 grad_norm 突然飆升到幾百甚至上千往往預(yù)示著即將發(fā)散這時(shí)候如果剛好有個 checkpoint就能回滾到健康狀態(tài)重來。total_norm torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) writer.add_scalar(train/grad_norm, total_norm, global_step) if total_norm 100: print(fWARNING: large grad norm {total_norm:.2f} at step {global_step})配合 TensorBoard 或 WandB把指標(biāo)實(shí)時(shí)可視化即使你不在服務(wù)器前也能通過手機(jī)看一眼訓(xùn)練是否正常。5.4 平臺配額與搶占的應(yīng)對策略云端 GPU 平臺經(jīng)常有配額限制或搶占機(jī)制。熱詞里提到的“配額已不夠預(yù)凍結(jié)”就是典型場景——你的任務(wù)可能因?yàn)榕漕~不足被暫停甚至回收。應(yīng)對策略是把 checkpoint 存到持久化存儲上而不是容器本地盤。容器一旦被回收本地盤的數(shù)據(jù)全沒了只有掛載的網(wǎng)絡(luò)存儲或?qū)ο蟠鎯δ鼙Ｗ?checkpoint。另外如果你的任務(wù)跑在搶占式實(shí)例上要假設(shè)它隨時(shí)可能被中斷。這種情況下 checkpoint 頻率要調(diào)高比如每 10 分鐘一次并且恢復(fù)邏輯要足夠健壯能在新實(shí)例上自動拉起。有些平臺提供了中斷通知信號比如收到 SIGTERM 后給你 30 秒善后可以注冊一個信號處理器在收到信號時(shí)立刻存一次 checkpointimport signal def handle_sigterm(signum, frame): print(Received SIGTERM, saving emergency checkpoint...) save_checkpoint(ckpt/emergency.pt, model, optimizer, scheduler, epoch, global_step, best_metric, scaler) sys.exit(0) signal.signal(signal.SIGTERM, handle_sigterm)這個“臨終存盤”機(jī)制在搶占式環(huán)境里價(jià)值極高能把損失從幾十分鐘壓縮到幾十秒。6. 一套可以直接抄的完整骨架把上面的東西串起來一個能扛住云端環(huán)境的訓(xùn)練腳本骨架大概長這樣import os, sys, time, signal, random import numpy as np import torch def main(): device torch.device(cuda if torch.cuda.is_available() else cpu) model build_model().to(device) optimizer build_optimizer(model) scheduler build_scheduler(optimizer) scaler torch.cuda.amp.GradScaler() ckpt_path find_latest_checkpoint(ckpt) if ckpt_path: start_epoch, global_step, best_metric load_checkpoint( ckpt_path, model, optimizer, scheduler, scaler) else: start_epoch, global_step, best_metric 0, 0, float(inf) def emergency_save(signum, frame): save_checkpoint(ckpt/emergency.pt, model, optimizer, scheduler, epoch, global_step, best_metric, scaler) sys.exit(0) signal.signal(signal.SIGTERM, emergency_save) last_save time.time() for epoch in range(start_epoch, NUM_EPOCHS): for batch in train_loader: loss train_step(model, batch, optimizer, scaler) global_step 1 if time.time() - last_save 1800: save_checkpoint(ckpt/last.pt, model, optimizer, scheduler, epoch, global_step, best_metric, scaler) last_save time.time() metric validate(model, val_loader) if metric best_metric: best_metric metric save_checkpoint(ckpt/best.pt, model, optimizer, scheduler, epoch, global_step, best_metric, scaler) if __name__ __main__: main()啟動命令放在 tmux 里外面套一層自動重啟的 shell 循環(huán)checkpoint 目錄掛到持久化存儲上。這套組合我在多個項(xiàng)目里用過從單卡微調(diào)到多卡預(yù)訓(xùn)練都扛得住最長的一次連續(xù)跑了十幾天沒出問題。最后分享一個我踩過的坑別把 checkpoint 和日志存在同一個目錄下用通配符清理。我曾經(jīng)寫過一個清理腳本rm ckpt/*.pt想刪舊 checkpoint結(jié)果手滑把best.pt也刪了而那次訓(xùn)練的最優(yōu)模型恰好只存在這一個文件里。從那以后我養(yǎng)成了習(xí)慣——最優(yōu)模型單獨(dú)放一個目錄清理腳本只針對last_*.pt這類滾動文件并且清理前先ls確認(rèn)一遍。這種低級錯誤往往比技術(shù)難題更讓人肉疼。