聊聊煉丹效率

煉丹多年，輾轉(zhuǎn)在不同地方待過，發(fā)現(xiàn)還是有相當(dāng)部分的小伙伴在手動(dòng)敲命令開所有的實(shí)驗(yàn)。高效一點(diǎn)的操作是寫一個(gè) bash script 然后用 for loop 把實(shí)驗(yàn)跑完，但似乎每次跑實(shí)驗(yàn)效率還是比較低，一輪下來也相當(dāng)累人。

總結(jié)下來有以下問題：

1. 效率。在迭代的早期通常會(huì)用小模型加小數(shù)據(jù)。單個(gè)模型往往只需占用一張卡，在常見的 8 卡服務(wù)器中便留著其他卡在那干等。
2. 認(rèn)知負(fù)載。訓(xùn)一個(gè)模型涉及到非常冗長(zhǎng)的 pipeline 以及眾多超參：數(shù)據(jù)處理，模型結(jié)構(gòu)，模型參數(shù)，訓(xùn)練參數(shù)，測(cè)試參數(shù)。這些可調(diào)的節(jié)點(diǎn)通常分布在代碼、數(shù)據(jù)或者命令行的參數(shù)里面。這讓檢查、排錯(cuò)和調(diào)整極其費(fèi)勁。往往開完一組實(shí)驗(yàn)一天的精力就見底了，更不用說隨時(shí)出現(xiàn)的錯(cuò)誤會(huì)讓這組實(shí)驗(yàn)的結(jié)果白白報(bào)廢。
3. 可用性。怎么在文件系統(tǒng)里面區(qū)分這些不同的實(shí)驗(yàn)？怎么樣高效地區(qū)分并分析這些實(shí)驗(yàn)的結(jié)果？怎么樣把在一個(gè) project 里面開發(fā)的工具快速遷移到其他的 project？
4. 魯棒性。如果機(jī)器突然宕機(jī)，哪些實(shí)驗(yàn)需要重新跑一遍？

筆者深受這些痛點(diǎn)折磨，所以一直也在尋找解決方案?，F(xiàn)在迭代的工作流感覺還算滿意，歡迎評(píng)論指教：

1. 把所有的模型、流程改動(dòng)都映射到命令行參數(shù)。模型結(jié)構(gòu)變化用 if/else 或者 switch/case 引出來。這符合把 operation 變成 code 的主流趨勢(shì)，也能夠無縫對(duì)接到大多數(shù)主流代碼框架。
2. 把不關(guān)心的默認(rèn)參數(shù)放到一個(gè)“命令模板”里面，然后將感興趣的參數(shù)變成變量，將感興趣的的取值組合寫到一個(gè)文件里面。把 baseline 取值也寫到文件里，方便對(duì)比。
3. 把文件當(dāng)做任務(wù)池，用一組 worker (分配了 gpu 資源) 并發(fā)地去拉任務(wù)，跑任務(wù)。每個(gè)任務(wù)的中間結(jié)果寫到以超參組合命名的文件夾中，這比時(shí)間戳更可讀，也有足夠區(qū)分不同實(shí)驗(yàn)，還可以查重防止重復(fù)跑實(shí)驗(yàn)。用 tensorboard 跟蹤訓(xùn)練進(jìn)程。
4. 寫一個(gè)評(píng)價(jià)標(biāo)準(zhǔn)和感興趣指標(biāo)的 parser 對(duì)所有實(shí)驗(yàn)的中間文件進(jìn)行處理，再把結(jié)果拉到 jupyter notebook 或者 excel 表里做可視化和分析。

為了降低認(rèn)知負(fù)擔(dān)，1 做了第一層簡(jiǎn)化：把分散在各個(gè)地方的調(diào)節(jié)點(diǎn)全部抽象成命令行參數(shù)，這樣只需在開發(fā)的時(shí)候保證每個(gè)調(diào)節(jié)點(diǎn)都正常 work 就行了，查錯(cuò)和決策就到了參數(shù)選擇的層面。2 則是做了第二層簡(jiǎn)化：將無關(guān)的參數(shù)和實(shí)驗(yàn)相關(guān)的參數(shù)區(qū)分開，這樣就避免了不小心改到別的參數(shù)而出錯(cuò) -- 而寫出這些無關(guān)的參數(shù)又迫使自己思考和過一遍所有可能的影響因子，查漏補(bǔ)缺。3 是實(shí)現(xiàn)層面的，任務(wù)池 + 并發(fā)的模式可以最大化硬件利用率，不需要惦記實(shí)驗(yàn)有沒有跑完。處理完這幾點(diǎn)后，基本上工作流就變成了白天讀 paper、溝通、debug、分析結(jié)果，晚上回家前列好要跑的實(shí)驗(yàn)跑起來，到家后看一下實(shí)驗(yàn)是否正常運(yùn)行，然后就可以倒頭睡覺等第二天出結(jié)果了。

基本原則講完之后也貼一個(gè)我的 python 實(shí)現(xiàn)。工具路徑在這里，希望大家實(shí)驗(yàn)跑的比誰都快：

https://github.com/simtony/runner

歡迎使用，star，二次開發(fā)，提 issue。非常相似的工具有微軟的 NNI（https://github.com/microsoft/nni），但是作為跑實(shí)驗(yàn)的工具來說太重了，而且做了太多抽象，很難對(duì)某個(gè)實(shí)驗(yàn)的結(jié)果做直觀的分析。另外這個(gè)回答（https://www.zhihu.com/question/384519338/answer/2152639948）提到的工具雖然實(shí)現(xiàn)了并發(fā)實(shí)驗(yàn)的功能，但感覺不太夠用。如果有更好的方案可以評(píng)論區(qū)分享一下。

舉個(gè)栗子

假設(shè)現(xiàn)在我們開發(fā)了一個(gè)新的 normalization 層叫 “newnorm”，baseline 是 batchnorm。每一個(gè)實(shí)驗(yàn)涉及 train、checkpoint average 和 test 三個(gè)流程?，F(xiàn)在希望看不同的 normalization 以及不同的 momentum 參數(shù)對(duì)結(jié)果的影響，對(duì)應(yīng)的配置文件如下：

---
template:
??train:?>
????python?train.py?data-bin/{data}
??????--seed?1
??????--criterion?label_smoothed_cross_entropy
??????--arch?transformer_iwslt_de_en?--share-all-embeddings
??????--optimizer?adam?--adam-betas?'(0.9,0.98)'?--clip-norm?0.0
??????--dropout?0.3?--lr-scheduler?inverse_sqrt?--warmup-updates?8000
??????--lr?0.0015?--min-lr?1e-09
??????--label-smoothing?0.1?--weight-decay?0.0001
??????--max-tokens?4096?
??????--save-dir?{_output}
??????--tensorboard-logdir?{_output}
??????--no-save-optimizer-state
??????--update-freq?1?--log-format?simple?--log-interval?50
??????--ddp-backend?no_c10d
??????--keep-last-epochs?5?--early-stop?5
??????--normalization?{norm}?[moment]

??avg:?>
????python?scripts/average_checkpoints.py?--inputs?{_output}
??????--num-epoch-checkpoints?5?--output?{_output}/averaged_model.pt

??test:?>
????python?generate.py?data-bin/{data}
????????--max-tokens?4096?--beam?5?--lenpen?1.0?--remove-bpe
????????--path?{_output}/averaged_model.pt?--gen-subset?test

default:
??data:?iwslt14
??norm:?batch
??moment:?0.1

resource:?[?0,?1,?2,?3?]

---
norm:?[?new,?batch?]
moment:?[?0.1,?0.05?]

第一個(gè) yaml doc 作為實(shí)驗(yàn)的 specification。template 下面指定了 train, checkpoint average 和 test 的模板命令，其中需要調(diào)的參數(shù)用 {param} 作為占位符。工具還定義了一些默認(rèn)的參數(shù)，比如這個(gè)實(shí)驗(yàn)對(duì)應(yīng)的路徑 {_output}。指定了要調(diào)的超參后，default 里面指定了這些超參的 baseline 值，最后在 resource 里指定了 4 個(gè) worker，每個(gè) worker 對(duì)應(yīng)一個(gè) GPU。

從第二個(gè) yaml doc 開始指定要格點(diǎn)搜的超參。默認(rèn)會(huì)把所有超參組合跑一遍。這里有 4 個(gè)任務(wù)。同步代碼和配置文件到服務(wù)器后，直接 run 并發(fā)地跑這4個(gè)任務(wù):

$ run
Orphan params: set()
Tasks: 4, commands to run: 12
START   gpu: 0, train: 1/ 4, output/Norm_new-Moment_0.1
START   gpu: 1, train: 2/ 4, output/Norm_new-Moment_0.05
START   gpu: 2, train: 3/ 4, output/Norm_batch-Moment_0.1
START   gpu: 3, train: 4/ 4, output/Norm_power-Moment_0.05
START   gpu: 2, avg  : 3/ 4, output/Norm_batch-Moment_0.1
FAIL    gpu: 2, avg  : 3/ 4, output/Norm_batch-Moment_0.1
...

每個(gè)輸出文件夾里面會(huì)寫入相應(yīng)的文件

$ ls output/Norm_batch-Moment_0.1
checkpoint51.pt
checkpoint52.pt
averaged_model.pt
log.train.20220316.030151
log.avg.20220316.030151
log.test.20220316.030151
param
stat

其中 log.* 是每個(gè)任務(wù)本來會(huì)打到命令行里面的 log。param 是每個(gè)任務(wù)對(duì)應(yīng)的一些參數(shù)設(shè)定，方便 debug，stat 則是任務(wù)狀態(tài)，分為success 和 fail。這可以用來幫助工具判斷是否需要重跑，也可以后期debug。跑實(shí)驗(yàn)的過程可以開 tensorboard 跟蹤結(jié)果，一旦不對(duì)勁馬上 kill。

實(shí)驗(yàn)跑完之后可以開一個(gè) jupyter notebook 寫實(shí)驗(yàn)結(jié)果分析的 parser。在這個(gè)例子里面只需要從 log.test 里面讀出 BLEU 就好了。寫完之后可以調(diào)用 Examiner 對(duì)所有結(jié)果做分析：

from?runner.examine?import?Examiner,?latest_log

#?define?a?metric?parser?for?each?directory?(experiment)
def?add_bleu(output_dir,?experiment,?caches):
????#?Each?parser?follows?the?same?signature
????#?It?can?read/write?to?a?global?cache?dict?`caches`,?
????#?and?read/write?each?experiment:?
????#?collections.namedtuple("Experiment",?["cache",?"metric",?"param"])
????latest_test_log?=?latest_log("test",?output_dir)
????bleu?=?parse_bleu(latest_test_log)?#??a?user-defined?log?parser
????experiment.metric["test_bleu"]?=?bleu
????
examiner?=?Examiner()??#?container?for?parsed?results
#?register?parser?for?each?directory?(experiment)
examiner.add(add_bleu)
#?run?all?parsers?for?directories?matched?by?regex?
examiner.exam(output="output",?regex=".*batch.*")
#?print?the?tsv?table?with?all?(different)?params?and?metrics?of?each?experiment
examiner.table()

為什么這樣寫

格點(diǎn)搜索可以適配大多數(shù)調(diào)參的場(chǎng)景。首先隨機(jī)暴力格點(diǎn)搜比較有效的調(diào)參方式，特別是當(dāng)計(jì)算資源比較充足的時(shí)候。其次做對(duì)比實(shí)驗(yàn)的時(shí)候也會(huì)用到參數(shù)的格點(diǎn)組合。最后如果不想格點(diǎn)搜，可以手動(dòng)把想跑的超參組合各自寫到配置文件里。

每一個(gè)實(shí)驗(yàn)都是由一系列順序執(zhí)行的命令組成的，比如上述例子的 train - checkpoint average - test。所以相比簡(jiǎn)單的命令，打包后的順序執(zhí)行的命令是更好的任務(wù)池的單元。

超參配置模式先后有兩個(gè)版本。一開始是直接定義一個(gè) config 類并對(duì)其操作。但是這樣會(huì)跟當(dāng)前 project 深度耦合，換一個(gè)代碼庫(kù)就得改很多地方，還會(huì)出錯(cuò)，并發(fā)部分也不好遷移到其他任務(wù)。最后將任務(wù)抽象成了一組命令，把修改超參轉(zhuǎn)化成修改任務(wù)命令，然后借用了 python 調(diào)用 bash 的接口進(jìn)行并發(fā)跑任務(wù)。這個(gè)方案完美匹配各大主流框架。

并發(fā)部分前后迭代了三個(gè)版本。第一版的 multiprocessing 最簡(jiǎn)單，但是對(duì)主進(jìn)程 Ctrl + C 后經(jīng)常出現(xiàn) orphan process，還需要查 pid 手動(dòng)去 kill。第二版的 thread 雖然沒有 orphan process 的問題，但是和 multiprocessing 一樣需要對(duì)全局共享的隊(duì)列和 io 加鎖，也很麻煩。最后收斂到了 asyncio 的 coroutine。后續(xù)加 cursor 的用戶界面也好寫一點(diǎn)。

一些不怎么高級(jí)的進(jìn)階功能

單個(gè) worker 需要多 GPU 的話可以在 resource 里面用引號(hào)框起來：resource: ["0,1", "2,3"] 。

日常需要一組實(shí)驗(yàn)在多個(gè)機(jī)器跑。不同機(jī)器卡數(shù)不同，需要跑的任務(wù)也不同。我先是用了 pycharm 的 deployment -> server group 的配置，讓每次 Ctrl + S都會(huì)把本地代碼 push 到所有服務(wù)器上。在上述工具方面做了幾個(gè)改動(dòng)：在命令行工具 run 中增加了 -t 和 -r。其中 -t 可以指定跑 yaml 文件中對(duì)應(yīng)_title 參數(shù)的任務(wù)。-r 指定 gpu index，這樣在不同機(jī)器通過命令行參數(shù)修改資源和 worker 數(shù)量。

經(jīng)常會(huì)出現(xiàn)跑一個(gè) train 和多個(gè) test 的情況。為了避免每次 test 都得從頭跑一次 train，加了 -c 命令來選擇要跑的 command。同時(shí)在 yaml 文件里面也加了 _cmd 字段方便按每組實(shí)驗(yàn)配置。

一個(gè)參數(shù)打包很多超參的情況也非常常見。典型的如 Transformer 的 pre/post layernorm 需要同時(shí)改 encoder 和 decoder 的 normalization 方式。在切換數(shù)據(jù)集的時(shí)候也是如此，不同數(shù)據(jù)往往意味著一整套超參的改變。所以在 yaml 的第一個(gè)文檔中加了 alias 字段，用來將某一個(gè)參數(shù)的取值映射到一組參數(shù)的取值。

有的時(shí)候快速試一些改進(jìn)的時(shí)候會(huì)懶得把它引到命令行參數(shù)里面。這個(gè)時(shí)候?yàn)榱藢?dāng)前結(jié)果和已有結(jié)果區(qū)分開，可以在參數(shù)選擇中引入 template 里面不用的參數(shù)。這些參數(shù)只會(huì)起到改輸出路徑名的作用。

公眾號(hào)后臺(tái)回復(fù)“CVPR 2022”獲取論文打包合集下載～