半夜調行車記錄器出來看,路面、車流、紅綠燈拍得一清二楚,偏偏整段畫面沒有一塊路牌。你知道車一直在開,就是說不出這是哪條路、哪個路口。看得到動作,對不上地圖。
說白了,今天要做的事就一件,幫你的 timeline 掛路牌。
昨天你學會把檔案瘦成三圈,打開一看,幾千顆 kernel 排排站,名字全是 ampere_sgemm_128x64_tn、vectorized_elementwise_kernel 這種天書。這些是 GPU 端的實作名,不是你的程式結構。哪段是 forward、哪段是 backward、哪顆是哪一層叫出來的,臉上通通沒寫。
但你其實看過掛好路牌的樣子。Day 4 那份 8 卡 Megatron 錄影,最上面那棵樹,Iteration、forward、backward 一層一層標得清清楚楚。那棵樹不是 nsys 送的,是 Megatron 的作者自己在程式裡標的。今天教你種同一棵。
工具叫 NVTX,PyTorch 內建就有包好。在 training loop 裡用 with 把每段包起來:
from torch.cuda import nvtx
for step, batch in enumerate(loader):
with nvtx.range(f"iter_{step}"):
with nvtx.range("data"):
batch = batch.cuda(non_blocking=True)
with nvtx.range("forward"):
loss = model(batch)
with nvtx.range("backward"):
loss.backward()
with nvtx.range("optim"):
optimizer.step(); optimizer.zero_grad()
nvtx.range 是 context manager,進 scope 推一個標記、出 scope 收掉,中間發生的一切都算這段的。它可以巢狀,iter_{step} 包著 data、forward、backward、optim 四塊,nsys 就會把它畫成一棵樹。名字是你自己取的,帶上 f-string 的 step 編號,timeline 上就能直接找到 iter_30 在哪。喜歡手動控制的也有 nvtx.range_push("名字") 跟 nvtx.range_pop() 這對,功能一樣,只是要自己負責成對。
除了範圍,NVTX 還有單點事件。nvtx.mark("ckpt saved") 在 timeline 上釘一根圖釘,適合標那種一瞬間的事,存檔完成、收到訊號、切換 phase 的瞬間。range 是一段,mark 是一點,兩個搭著用。
錄影指令一個字都不用改。Day 6 那行 --trace=cuda,nvtx 本來就把 nvtx 的耳朵打開了,你標了它就錄,沒標那條列就空著。而且 Day 6 的冷知識在這裡再用一次,NVTX 沒 profiler 掛著時是空殼呼叫,這些 with 常駐在 code 裡零成本,跟昨天的 cudaProfilerStart() 一樣,標了就不用拔。
動手之前還可以先撿現成的。不少框架本來就幫你標好了,Day 4 那份 Megatron 錄影的樹就是官方自己種的。所以拿到一個新專案,先照 Day 6 錄一小段開來看,NVTX 列上已經有樹,你就省工了;空的,再回來自己標。

先讓你少走一段冤枉路。NVTX 記錄的時間點,是 CPU 執行到 push 跟 pop 那兩行的瞬間。可是 Day 4 就講過,訂單下在 CPU、出貨在 GPU,kernel 是非同步的,晚點才真的跑。
所以你的 forward range 蓋住的,是 CPU 下單 forward 那段時間,不是 GPU 真的在算 forward 的時間。這兩段可以差很遠,CPU 常常早就跑去下 backward 的單了,GPU 還在出 forward 的貨。你要是直接拿 CPU 側那條 range 的長度當 forward 的耗時,量到的是下單速度,不是出貨速度,又繞回 Day 3 那種指標錯位的老問題。
好消息是 Nsight 有處理這件事。它會把 NVTX range 投影(project)到 GPU 側,做法是看這段 range 裡發射出去的 kernel 實際在 GPU 上佔了哪段時間,重畫一條 GPU 版的 NVTX 列。nsys stats 的報表家族裡也有一支 nvtx_gpu_proj_trace 專門吐這個。判準記一句,要回答 forward 在 GPU 上花多久,看投影後的,不是 CPU 那條原始 range。
樹種好、重錄一份,打開 timeline,教你先讀兩個現象。
第一個,forward 跟 backward 的比例。有個底可以先放在心裡,backward 通常吃 forward 兩倍上下的時間,機制很單純,forward 一層做一次矩陣乘,backward 要算兩份,一份對 input 的梯度傳給上一層、一份對 weight 的梯度給 optimizer 用,帳面上就是兩次。所以看到 backward 大約是 forward 的兩倍,正常,別急著修。反過來,比例明顯超過兩倍,常見的嫌犯是 activation checkpointing,為了省記憶體,backward 裡會把 forward 重算一遍,比例自然被推向三倍。這時候那不是病,是你自己選的以時間換空間,看到帳要認得出來。
第二個,圈與圈之間有沒有縫。這裡有個細節要看懂,batch 是在 for 那一行從 loader 拿出來的,發生在 iter range 打開之前,所以 dataloader 等資料的時間不會被任何牌子蓋到,它會現形成 iter_30 結束到 iter_31 開始之間的一段空白。看到圈與圈之間規律地空一截,就是 loader 餵不上;data 段本身肥得不像話,則是搬上 GPU 的那段在拖。兩個都是 Day 2 死法二的臉,Day 16 會專門修。以前這種問題要瞇著眼睛猜,現在路牌一掛,沒被牌子蓋到的縫反而最顯眼。
再送一個加碼的,瞄一眼 optim 段長什麼樣。它通常不長,但你放大看常會嚇一跳,裡面是密密麻麻一大票小 kernel,傳統實作是每個參數 tensor 更新一次就發射一顆(新版 PyTorch 的 foreach 跟 fused optimizer 會把它們併起來,所以你看到的密度跟版本有關)。段不長沒事,要是 optim 段意外地肥、而且全是幾 µs 的小不點,那是 launch 次數太多在吃你,Day 15 講 launch-bound 的時候它會再出場。
手動標 phase 夠用九成的場合,但偶爾你想看更細,細到每一個算子。PyTorch 有內建的自動標:
with torch.autograd.profiler.emit_nvtx():
# 跑你要錄的那幾步
emit_nvtx 會給每個 aten 算子自動掛一個 range,而且 forward 的 range 上會寫 seq=N,backward 對應的那段寫 stashed seq=N,兩個號碼一對,你就能把 backward 裡慢的那段,對回它是 forward 哪個 op 生出來的。當 backward 慢得不像話、遠超過該有的比例,要揪出裡面哪段在搞鬼的時候,這是神器。加上 record_shapes=True 還會把 tensor 的 shape 一起寫進 range 名,連是哪個尺寸的矩陣乘都告訴你。
代價也很直白。一步幾百上千個 range,timeline 吵得要命,overhead 也跟著漲,這又是 Day 6 那筆帳,事件越多、干擾越大。所以分工是,手動標的 phase 層級當常駐路牌,emit_nvtx 當偶爾才開的逐字稿模式,確認完就關。一個是章節目錄,一個是逐字稿,你不會用逐字稿當目錄。
NVTX 不是 PyTorch 的專利。pip install nvtx 之後,任何 Python code 都能標,nvtx.annotate("load_data") 當 context manager 或 decorator 都行,dataloader、前處理、自己寫的資料清洗通通標得到。C++ 或 CUDA 端則是 nvtxRangePush / nvtxRangePop 這對原生 API。你的 pipeline 有幾層,路牌就能掛到幾層。
再往下挖一層,有幾件 NVTX 的內行事,早知道早省事。
第一件,push/pop 是一疊 thread-local 的疊疊樂。每條 thread 自己維護自己的一疊,你在哪條 thread 上 push,range 就畫在 timeline 上那條 thread 的列,巢狀關係也只在同一條 thread 裡成立。這解釋了一個新手常見的困惑,在主行程種的樹,為什麼看不到別條 thread 在做的事,因為人家的帳記在人家的列上。
第二件,跨 thread 的區間有另一對 API。push/pop 管不了"這頭開、那頭關"的場景,比如一個請求從收到(thread A)到回完(thread B)。NVTX 另有 start/end 款,開的時候拿到一個 id,任何 thread 拿這個 id 都能關,pip 版的 nvtx.start_range() / nvtx.end_range(id) 就是它,區間可以交錯、不必巢狀。做推論服務想標"一個 request 的一生",用這對才標得起來。
第三件,接著上面,DataLoader 的 worker 根本不是 thread,是 fork 出去的另一個 process。你在 worker 函式裡標的 range,會畫在那個 worker process 自己的列上,不會長進主行程的樹裡。這不是 bug,是行程的邊界,明天 Day 9 講 fork 錄影的坑時,這件事會再回來咬一口。
第四件,路牌可以上色。NVTX 的規格本身支援給 range 指定顏色跟數值 payload,pip 版寫 nvtx.annotate("data", color="orange") 就行,在 GUI 裡掃 timeline 的時候,顏色比字快得多。torch 包的那層只讓你給名字,要上色就改用 pip 版的 nvtx,兩邊可以混用。
第五件,別標錯耳朵。PyTorch 自家的 torch.profiler 聽的是 record_function 這套標記,nsys 聽的是 NVTX,兩邊互不相通。你用 nsys 錄,就標 NVTX;用 torch.profiler 出 chrome trace,才用 record_function。標錯邊,工具一個字都聽不到,白標。
掛好路牌,前幾天學的每一招都跟著升級。
Day 5 的報表家族多一支可用,nsys stats --report nvtx_sum 會把你的 range 攤成一張表(示意):
Time(%) Total Time Instances Range
48.1 1.42 s 3 backward
26.7 0.79 s 3 forward
12.4 0.37 s 3 optim
8.9 0.26 s 3 data
每個名字出現幾次、總共吃掉多久、佔比多少,一行指令就有數字,而且這次表上寫的是你自己取的名字,不是天書。
這張表還藏著一個命名的眉角。它是按名字聚合的,同名的湊成一列。所以 forward 這種固定名字,三圈會漂亮地加總成一列;但 iter_30、iter_31 這種帶編號的,每圈名字都不同,表上會散成一人一列,聚不起來(上面那張示意表,我就偷偷把 iter 那三列藏掉了,不然它們會各佔一行)。判準是,想在報表上聚合的層級用固定名,想在 timeline 上定位的層級才帶編號,兩個需求別塞在同一個名字裡。
昨天埋的鉤子也正式解鎖。-c 填 nvtx、-p 給一個 range 名,nsys 就拿你的標記當開錄的開關:
nsys profile --trace=cuda,nvtx -c nvtx -p "iter_300" -o tagged python train.py
程式跑到名叫 iter_300 的 range 才開錄,range 結束就收,連 cudaProfilerStart() 都不用喊。注意它圈的是一段 range,所以上面這行錄到的是一圈;想一次圈三圈,就另外標一個橫跨那三步的 range(比如叫 steady),-p 指它就好。圈在哪裡跟那裡叫什麼名字,從今天起是同一件事。
順帶補一個 -p 語法裡的小字,它完整的格式其實是 range@domain。domain 是 NVTX 的命名空間,函式庫可以把自己的標記收在自己的 domain 裡,跟你的分開,撞名也不怕。平常自己標的都落在預設 domain,寫 range 名就好;哪天要指定某個框架自帶的標記當開關,才需要 @ 後面那截。
最後一個判準,路牌跟著問題走,不是跟著虔誠走。
想回答哪個 phase 慢,標到 phase 這層就夠。懷疑某幾層有鬼,再對那幾層加標。一開始就把模型每一層都包一圈 range 的人,跟昨天錄好錄滿的人是同一種人,timeline 會被幾百段套疊的字幕淹掉,你反而什麼都看不到。路牌是幫人認路的,不是把整條路貼滿的。
我自己的起手式就是開頭那段 code 的長相,iter 裡切四塊,data、forward、backward、optim。這四塊就足夠回答八成的第一個問題,剩下兩成,等證據指到哪裡再往哪裡加標。
真的走到要一層層標的那天,也不用手改模型的每一層。PyTorch 的 module hook 可以自動包,register_forward_pre_hook 進場時 push、register_forward_hook 出場時 pop,對著想看的那幾個 module 掛上去,路牌就自己長出來,看完摘掉 hook 就還原。工具是死的,掛牌的顆粒度永遠掌握在你手上。
今天 timeline 會講話了。五行 nvtx.range 種出一棵樹、emit_nvtx 對算子、投影看 GPU 側的真實耗時、nvtx_sum 出表、-c nvtx -p 拿標記當錄影開關。從此你看到的不是幾千顆匿名 kernel,是你自己程式的章節目錄。
不過到目前為止,我們都假設你在一台乾淨的機器上直接跑 python train.py。真實世界不長這樣,你的 job 八成關在 Docker 容器裡、排在 Slurm 隊伍裡,在那裡面按錄影鍵,會遇到一整包新的坑,權限不夠、工具不在、fork 出來的行程錄不到。明天 Day 9,我們把行車記錄器裝進貨櫃裡。
路牌掛好了。明天,開進港口。
torch.cuda.nvtx.range/range_push/range_pop:PyTorch docs — torch.cuda.nvtx。emit_nvtx 與 seq/stashed seq 對應機制:PyTorch docs — torch.autograd.profiler.emit_nvtx。nvtx_gpu_proj_trace)與 nvtx_sum:Nsight Systems User Guide — stats reports。