煉丹開掛指南！！！

來源丨h(huán)ttps://zhuanlan.zhihu.com/p/482876481；作者：歐澤彬，西湖大學(xué)；

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

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

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

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

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

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

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

https://github.com/simtony/runner

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

舉個栗子

假設(shè)現(xiàn)在我們開發(fā)了一個新的 normalization 層叫 “newnorm”，baseline 是 batchnorm。每一個實驗涉及 train、checkpoint average 和 test 三個流程。現(xiàn)在希望看不同的 normalization 以及不同的 momentum 參數(shù)對結(jié)果的影響，對應(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?]

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

從第二個 yaml doc 開始指定要格點搜的超參。默認會把所有超參組合跑一遍。這里有 4 個任務(wù)。同步代碼和配置文件到服務(wù)器后，直接?run?并發(fā)地跑這4個任務(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
...

每個輸出文件夾里面會寫入相應(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.*?是每個任務(wù)本來會打到命令行里面的 log。param?是每個任務(wù)對應(yīng)的一些參數(shù)設(shè)定，方便 debug，stat?則是任務(wù)狀態(tài)，分為success?和?fail。這可以用來幫助工具判斷是否需要重跑，也可以后期debug。跑實驗的過程可以開 tensorboard 跟蹤結(jié)果，一旦不對勁馬上 kill。

實驗跑完之后可以開一個 jupyter notebook 寫實驗結(jié)果分析的 parser。在這個例子里面只需要從?log.test?里面讀出 BLEU 就好了。寫完之后可以調(diào)用?Examiner?對所有結(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()

為什么這樣寫

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

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

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

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

一些不怎么高級的進階功能

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

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

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

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