大模型訓練的GPU聯手CPU顯存優化分析方法
本文經自動駕駛之心公眾號授權轉載，轉載請聯系出處。
隨著深度學習模型復雜度和數據集規模的增大，計算效率成為了不可忽視的問題。GPU憑借強大的并行計算能力，成為深度學習加速的標配。然而，由于服務器的顯存非常有限，隨著訓練樣本越來越大，顯存連一個樣本都容不下的現象頻頻發生。除了升級硬件（燒錢）、使用分布式訓練（費力），你知道還有哪些方法嗎？即使顯存充足，所有運算都在GPU上執行就是最高效嗎？只要掌握以下小知識，模型訓練的種種問題統統搞定，省時省力省錢，重點是高效！
其實CPU和GPU是協同工作的，如果能合理地利用它們各自的優勢，就能夠節省顯存資源（顯存不夠內存來湊），甚至獲得更好的訓練性能。本文為您提供了device_guard接口，只需要一行命令，即可實現GPU和CPU的混合訓練，不僅可以解決訓練模型時通過調整批尺寸（batch size）顯存依然超出的問題，讓原本無法在單臺服務器執行的模型可以訓練，同時本文還給出了提高GPU和CPU混合訓練效率的方法，將服務器資源利用到極致，幫助您提升模型的性能！
一、模型訓練的特點
深度學習任務通常使用GPU進行模型訓練。這是因為GPU相對于CPU具有更多的算術邏輯單元（ALU），可以發揮并行計算的優勢，特別適合計算密集型任務，可以更高效地完成深度學習模型的訓練。GPU模式下的模型訓練如圖1所示，總體可以分為4步：
第1步，將輸入數據從系統內存拷貝到顯存。
第2步，CPU指示GPU處理數據。
第3步，GPU并行地完成一系列的計算。
第4步，將計算結果從顯存拷貝到內存。
圖1 模型訓練示意圖
從圖中可以了解到，雖然GPU并行計算能力優異，但無法單獨工作，必須由CPU進行控制調用；而且顯存和內存之間的頻繁數據拷貝也可能會帶來較大的性能開銷。CPU雖然計算能力不如GPU，但可以獨立工作，可以直接訪問內存數據完成計算。因此，想獲得更好的訓練性能，需要合理利用GPU和CPU的優勢。
二、模型訓練的常見問題
問題一：GPU顯存爆滿，資源不足
你建的模型不錯，在這個簡潔的任務中可能成為新的SOTA，但每次嘗試批量處理更多樣本時，你都會得到一個CUDA RuntimeError：out of memory。
這是因為GPU卡的顯存是非常有限的，一般遠低于系統內存。以V100為例，其顯存最高也僅有32G，甚至有些顯存僅12G左右。因此當模型的參數量較大時，在GPU模式下模型可能無法訓練起來。
設置CPU模式進行模型訓練，可以避免顯存不足的問題，但是訓練速度往往太慢。
那么有沒有一種方法，可以在單機訓練中充分地利用GPU和CPU資源，讓部分層在CPU執行，部分層在GPU執行呢？
問題二：頻繁數據拷貝，訓練效率低
在顯存足夠的情況下，我們可以直接采用GPU模式去訓練模型，但是讓所有的網絡層都運行在GPU上就一定最高效嗎？其實GPU只對特定任務更快，而CPU擅長各種復雜的邏輯運算。框架中有一些OP會默認在CPU上執行，或者有一些OP的輸出會被存儲在CPU上，因為這些輸出往往需要在CPU上訪問。這就會導致訓練過程中，CPU和GPU之間存在數據拷貝。
圖2是CPU和GPU數據傳輸示意圖。假設模型的中間層存在下圖中的4個算子。其中算子A和算子B都在CPU執行，因此B可以直接使用A的輸出。算子C和算子D都在GPU上執行，那么算子D也可以直接使用C的輸出。但是算子B執行完，其輸出在CPU上，在算子C執行時，就會將B的輸出從CPU拷貝到GPU。
頻繁的數據拷貝，也會影響模型的整體性能。如果能把算子A和B設置在GPU上執行，或者算子C和D設置在CPU上執行，避免數據傳輸，或許會提升模型性能。那么應該如何更加合理地為算子分配設備，使得訓練過程更加高效呢？我們需要更綜合地考慮，在發揮GPU和CPU各自計算優勢的前提下，降低數據拷貝帶來的時間消耗。
圖2 CPU和GPU數據傳輸示意圖
三、定義化GPU和CPU混合訓練
上面兩個場景都是希望為模型中的部分層指定運行設備。飛槳提供了fluid.CUDAPlace和fluid.CPUPlace用于指定運行設備，但這兩個接口在指定設備時是二選一的，也就是說要么在GPU模式下訓練，要么在CPU模式下訓練。過去我們無法指定某一部分計算在GPU上執行還是在CPU上執行。飛槳開源框架從1.8版本開始提供了device_guard接口，使用該接口可以為網絡中的計算層指定設備為CPU或者GPU，實現更靈活的異構計算調度。
如何使用device_guard接口解決上面兩個場景中提到的問題呢？接下來，我們看看具體的例子。
好處一：充分利用CPU資源，避免顯存超出
如果使用fluid.CUDAPlace指定了全局的運行設備，飛槳將會自動把支持GPU計算的OP分配在GPU上執行，然而當模型參數量過大并且顯存有限時，很可能會遇到顯存超出的情況。如下面的示例代碼，embedding層的參數size包含兩個元素，第一個元素為vocab_size(詞表大小)，第二個為emb_size（embedding層維度）。實際場景中，詞表可能會非常大。示例代碼中，詞表大小被設置為10,000,000，該層創建的權重矩陣的大小為(10000000, 150)，僅這一層就需要占用5.59G的顯存。如果再加上其他的網絡層，在這種大詞表場景下，很有可能會顯存超出。
import paddle.fluid as fluid
data = fluid.layers.fill_constant(shape=[1], value=128, dtype='int64')
label = fluid.layers.fill_constant(shape=[1, 150], value=0.5, dtype='float32')
emb = fluid.embedding(input=data, size=(10000000, 150), dtype='float32')
out = fluid.layers.l2_normalize(x=emb, axis=-1)
cost = fluid.layers.square_error_cost(input=out, label=label)
avg_cost = fluid.layers.mean(cost)
sgd_optimizer = fluid.optimizer.SGD(learning_rate=0.001)
sgd_optimizer.minimize(avg_cost)
place = fluid.CUDAPlace(0)
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
result = exe.run(fluid.default_main_program(), fetch_list=[avg_cost])embedding是根據input中的id信息從embedding矩陣中查詢對應embedding信息，它并不是一個計算密度非常高的OP，因此在CPU上進行計算，其速度也是可接受的。如果將embedding層設置在CPU上運行，就能夠充分利用CPU大內存的優勢，避免顯存超出。可以參考如下代碼，使用device_guard將embedding層設置在CPU上。那么，除了embedding層，其他各層都會在GPU上運行。
import paddle.fluid as fluid
data = fluid.layers.fill_constant(shape=[1], value=128, dtype='int64')
label = fluid.layers.fill_constant(shape=[1, 150], value=0.5, dtype='float32')
with fluid.device_guard("cpu"): #一行命令，指定該網絡層運行設備為CPU
emb = fluid.embedding(input=data, size=(10000000, 150), dtype='float32')
out = fluid.layers.l2_normalize(x=emb, axis=-1)
cost = fluid.layers.square_error_cost(input=out, label=label)
avg_cost = fluid.layers.mean(cost)
sgd_optimizer = fluid.optimizer.SGD(learning_rate=0.001)
sgd_optimizer.minimize(avg_cost)
place = fluid.CUDAPlace(0)
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
result = exe.run(fluid.default_main_program(), fetch_list=[avg_cost])因此，在顯存有限時你可以參考上面的示例將一些計算密度不高的網絡層設置在CPU上避免顯存超出。
好處二：合理設置運行設備，減少數據傳輸
如果你在GPU模式下訓練模型，希望提升訓練速度，那么可以看看模型中是否存在一些不必要的數據傳輸。在文章開頭我們提到CPU和GPU之間的數據拷貝是耗時的，因此如果能夠避免這樣的情況，就有可能提升模型的性能。
在下面的內容中，我們將教你如何通過profile工具分析數據傳輸開銷，以及如何使用device_guard避免不必要的數據傳輸，從而提升模型性能。大致流程如下：
- 首先使用profile工具對模型進行分析，查看是否存在GpuMemcpySync的調用耗時。若存在，則進一步分析發生數據傳輸的原因。
- 通過Profiling Report找到發生GpuMemcpySync的OP。如果需要，可以通過打印log，找到GpuMemcpySync發生的具體位置。
- 嘗試使用device_guard設置部分OP的運行設備，來減少GpuMemcpySync的調用。
- 最后比較修改前后模型的Profiling Report，或者其他用來衡量性能的指標，確認修改后是否帶來了性能提升。
步驟1、使用profile工具確認是否發生了數據傳輸
首先我們需要分析模型中是否存在CPU和GPU之間的數據傳輸。在OP執行過程中，如果輸入Tensor所在的設備與OP執行的設備不同，就會自動將輸入Tensor從CPU拷貝到GPU，或者從GPU拷貝到CPU，這個過程是同步的數據拷貝，通常比較耗時。下列示例代碼的14行設置了profile，利用profile工具我們可以看到模型的性能數據。
import paddle.fluid as fluid
import paddle.fluid.compiler as compiler
import paddle.fluid.profiler as profiler
data1 = fluid.layers.fill_constant(shape=[1, 3, 8, 8], value=0.5, dtype='float32')
data2 = fluid.layers.fill_constant(shape=[1, 3, 5, 5], value=0.5, dtype='float32')
shape = fluid.layers.shape(data2)
shape = fluid.layers.slice(shape, axes=[0], starts=[0], ends=[4])
out = fluid.layers.crop_tensor(data1, shape=shape)
place = fluid.CUDAPlace(0)
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
compiled_prog = compiler.CompiledProgram(fluid.default_main_program())
with profiler.profiler('All', 'total') as prof:
for i in range(10):
result = exe.run(program=compiled_prog, fetch_list=[out])在上述程序運行結束后，將會自動地打印出下面的Profiling Report，可以看到GpuMemCpy Summary中給出了2項數據傳輸的調用耗時。如果GpuMemCpy Summary中存在GpuMemcpySync，那么就說明你的模型中存在同步的數據拷貝。
進一步分析，可以看到slice和crop_tensor執行中都發生了GpuMemcpySync。我們通過查看網絡的定義，就會發現盡管我們在程序中設置了GPU模式運行，但是shape這個OP將輸出結果存放在CPU上，導致后面在GPU上執行的slice使用這個結果時發生了從CPU到GPU的數據拷貝。slice的輸出結果存放在GPU上，而crop_tensor用到這個結果的參數默認是從CPU上取數據，因此又發生了一次數據拷貝。
-------------------------> Profiling Report <-------------------------
Note! This Report merge all thread info into one.
Place: All
Time unit: ms
Sorted by total time in descending order in the same thread
Total time: 26.6328
Computation time Total: 13.3133 Ratio: 49.9884%
Framework overhead Total: 13.3195 Ratio: 50.0116%
------------------------- GpuMemCpy Summary -------------------------
GpuMemcpy Calls: 30 Total: 1.47508 Ratio: 5.5386%
GpuMemcpyAsync Calls: 10 Total: 0.443514 Ratio: 1.66529%
GpuMemcpySync Calls: 20 Total: 1.03157 Ratio: 3.87331%
------------------------- Event Summary -------------------------
Event Calls Total CPU Time (Ratio) GPU Time (Ratio) Min. Max. Ave. Ratio.
FastThreadedSSAGraphExecutorPrepare 10 9.16493 9.152509 (0.998645) 0.012417 (0.001355) 0.025192 8.85968 0.916493 0.344122
shape 10 8.33057 8.330568 (1.000000) 0.000000 (0.000000) 0.030711 7.99849 0.833057 0.312793
fill_constant 20 4.06097 4.024522 (0.991025) 0.036449 (0.008975) 0.075087 0.888959 0.203049 0.15248
slice 10 1.78033 1.750439 (0.983212) 0.029888 (0.016788) 0.148503 0.290851 0.178033 0.0668471
GpuMemcpySync:CPU->GPU 10 0.45524 0.446312 (0.980388) 0.008928 (0.019612) 0.039089 0.060694 0.045524 0.0170932
crop_tensor 10 1.67658 1.620542 (0.966578) 0.056034 (0.033422) 0.143906 0.258776 0.167658 0.0629515
GpuMemcpySync:GPU->CPU 10 0.57633 0.552906 (0.959357) 0.023424 (0.040643) 0.050657 0.076322 0.057633 0.0216398
Fetch 10 0.919361 0.895201 (0.973721) 0.024160 (0.026279) 0.082935 0.138122 0.0919361 0.0345199
GpuMemcpyAsync:GPU->CPU 10 0.443514 0.419354 (0.945526) 0.024160 (0.054474) 0.040639 0.059673 0.0443514 0.0166529
ScopeBufferedMonitor::post_local_exec_scopes_process 10 0.341999 0.341999 (1.000000) 0.000000 (0.000000) 0.028436 0.057134 0.0341999 0.0128413
eager_deletion 30 0.287236 0.287236 (1.000000) 0.000000 (0.000000) 0.005452 0.022696 0.00957453 0.010785
ScopeBufferedMonitor::pre_local_exec_scopes_process 10 0.047864 0.047864 (1.000000) 0.000000 (0.000000) 0.003668 0.011592 0.0047864 0.00179718
InitLocalVars 1 0.022981 0.022981 (1.000000) 0.000000 (0.000000) 0.022981 0.022981 0.022981 0.000862883步驟2、通過log查看發生數據傳輸的具體位置
有時同一個OP會在模型中被用到很多次，例如可能我們會在網絡的幾個不同位置，都定義了slice層。這時候想要確認究竟是在哪個位置發生了數據傳輸，就需要去查看更加詳細的調試信息，那么可以打印出運行時的log。依然以上述程序為例，執行GLOG_vmodule=operator=3 python test_case.py，會得到如下log信息，可以看到這兩次數據傳輸：
- 第3～7行log顯示：shape輸出的結果在CPU上，在slice運行時，shape的輸出被拷貝到GPU上
- 第9～10行log顯示：slice執行完的結果在GPU上，當crop_tensor執行時，它會被拷貝到CPU上。
I0406 14:56:23.286592 17516 operator.cc:180] CUDAPlace(0) Op(shape), inputs:{Input[fill_constant_1.tmp_0:float[1, 3, 5, 5]({})]}, outputs:{Out[shape_0.tmp_0:int[4]({})]}.
I0406 14:56:23.286628 17516 eager_deletion_op_handle.cc:107] Erase variable fill_constant_1.tmp_0 on CUDAPlace(0)
I0406 14:56:23.286725 17516 operator.cc:1210] Transform Variable shape_0.tmp_0 from data_type[int]:data_layout[NCHW]:place[CPUPlace]:library_type[PLAIN] to data_type[int]:data_layout[ANY_LAYOUT]:place[CUDAPlace(0)]:library_type[PLAIN]
I0406 14:56:23.286763 17516 scope.cc:169] Create variable shape_0.tmp_0
I0406 14:56:23.286784 17516 data_device_transform.cc:21] DeviceTransform in, src_place CPUPlace dst_place: CUDAPlace(0)
I0406 14:56:23.286867 17516 tensor_util.cu:129] TensorCopySync 4 from CPUPlace to CUDAPlace(0)
I0406 14:56:23.287099 17516 operator.cc:180] CUDAPlace(0) Op(slice), inputs:{EndsTensor[], EndsTensorList[], Input[shape_0.tmp_0:int[4]({})], StartsTensor[], StartsTensorList[]}, outputs:{Out[slice_0.tmp_0:int[4]({})]}.
I0406 14:56:23.287140 17516 eager_deletion_op_handle.cc:107] Erase variable shape_0.tmp_0 on CUDAPlace(0)
I0406 14:56:23.287220 17516 tensor_util.cu:129] TensorCopySync 4 from CUDAPlace(0) to CPUPlace
I0406 14:56:23.287473 17516 operator.cc:180] CUDAPlace(0) Op(crop_tensor), inputs:{Offsets[], OffsetsTensor[], Shape[slice_0.tmp_0:int[4]({})], ShapeTensor[], X[fill_constant_0.tmp_0:float[1, 3, 8, 8]({})]}, outputs:{Out[crop_tensor_0.tmp_0:float[1, 3, 5, 5]({})]}.步驟3、使用device_guard避免不必要的數據傳輸
在上面的例子中，shape輸出的是一個1-D的Tensor，因此在slice執行時，計算代價相對于數據傳輸代價或許是更小的。如果將slice設置在CPU上運行，就可以避免2次數據傳輸，那么是不是有可能提升模型速度呢？我們嘗試修改程序，將slice層設置在CPU上執行：
import paddle.fluid as fluid
import paddle.fluid.compiler as compiler
import paddle.fluid.profiler as profiler
data1 = fluid.layers.fill_constant(shape=[1, 3, 8, 8], value=0.5, dtype='float32')
data2 = fluid.layers.fill_constant(shape=[1, 3, 5, 5], value=0.5, dtype='float32')
shape = fluid.layers.shape(data2)
with fluid.device_guard("cpu"): # 一行命令，指定該網絡層運行設備為CPU
shape = fluid.layers.slice(shape, axes=[0], starts=[0], ends=[4])
out = fluid.layers.crop_tensor(data1, shape=shape)
place = fluid.CUDAPlace(0)
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
compiled_prog = compiler.CompiledProgram(fluid.default_main_program())
with profiler.profiler('All', 'total') as prof:
for i in range(10):
result = exe.run(program=compiled_prog, fetch_list=[out])步驟4、比較修改前后模型，確認是否帶來性能提升
再次觀察Profiling Report 中GpuMemCpy Summary的內容，可以看到GpuMemCpySync 這一項已經被消除了。同時注意到，下面的Total time為 14.5345 ms，而修改前是26.6328 ms，速度提升一倍！此實驗說明使用device_guard避免數據傳輸后，示例模型的性能有了明顯的提升。 在實際的模型中，若GpuMemCpySync調用耗時占比較大，并且可以通過設置device_guard避免，那么就能夠帶來一定的性能提升。
-------------------------> Profiling Report <-------------------------
Note! This Report merge all thread info into one.
Place: All
Time unit: ms
Sorted by total time in descending order in the same thread
Total time: 14.5345
Computation time Total: 4.47587 Ratio: 30.7948%
Framework overhead Total: 10.0586 Ratio: 69.2052%
------------------------- GpuMemCpy Summary -------------------------
GpuMemcpy Calls: 10 Total: 0.457033 Ratio: 3.14447%
GpuMemcpyAsync Calls: 10 Total: 0.457033 Ratio: 3.14447%
------------------------- Event Summary -------------------------
Event Calls Total CPU Time (Ratio) GPU Time (Ratio) Min. Max. Ave. Ratio.
FastThreadedSSAGraphExecutorPrepare 10 7.70113 7.689066 (0.998433) 0.012064 (0.001567) 0.032657 7.39363 0.770113 0.529852
fill_constant 20 2.62299 2.587022 (0.986287) 0.035968 (0.013713) 0.071097 0.342082 0.13115 0.180466
shape 10 1.93504 1.935040 (1.000000) 0.000000 (0.000000) 0.026774 1.6016 0.193504 0.133134
Fetch 10 0.880496 0.858512 (0.975032) 0.021984 (0.024968) 0.07392 0.140896 0.0880496 0.0605797
GpuMemcpyAsync:GPU->CPU 10 0.457033 0.435049 (0.951898) 0.021984 (0.048102) 0.037836 0.071424 0.0457033 0.0314447
crop_tensor 10 0.705426 0.671506 (0.951916) 0.033920 (0.048084) 0.05841 0.123901 0.0705426 0.0485346
slice 10 0.324241 0.324241 (1.000000) 0.000000 (0.000000) 0.024299 0.07213 0.0324241 0.0223084
eager_deletion 30 0.250524 0.250524 (1.000000) 0.000000 (0.000000) 0.004171 0.016235 0.0083508 0.0172365
ScopeBufferedMonitor::post_local_exec_scopes_process 10 0.047794 0.047794 (1.000000) 0.000000 (0.000000) 0.003344 0.014131 0.0047794 0.00328831
InitLocalVars 1 0.034629 0.034629 (1.000000) 0.000000 (0.000000) 0.034629 0.034629 0.034629 0.00238254
ScopeBufferedMonitor::pre_local_exec_scopes_process 10 0.032231 0.032231 (1.000000) 0.000000 (0.000000) 0.002952 0.004076 0.0032231 0.00221755總結
通過以上實驗對比可以發現，device_guard接口能夠做到一條命令即可合理設置模型網絡層的運行設備，對模型進行GPU和CPU計算的更靈活調度，將服務器的資源利用到極致，解決顯存容量捉襟見肘導致模型無法訓練的問題。怎么樣，這個功能是不是相當實用！心動不如心動，快快參考本文的方法，盡情訓練自己的模型吧！
