訓(xùn)練實(shí)戰(zhàn)：從零訓(xùn)練 RoBERTa、GPT-2、T5 與 BART)
推理引擎大模型【免費(fèi)下載鏈接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.項(xiàng)目地址https://gitcode.com/gh_mirrors/fl/FlexGen點(diǎn)擊查看免費(fèi)下載本篇技術(shù)指南以benchmark/third_party/transformers/examples/flax/language-modeling/README.md為核心文檔系統(tǒng)講解如何利用 JAX/Flax 后端從零預(yù)訓(xùn)練四種主流語言模型——雙向的 RoBERTaMLM、自回歸的 GPT-2CLM、span-masked 的 T5 以及去噪自編碼的 BART。文中的示例腳本與配套代碼均位于當(dāng)前倉庫的benchmark/third_party/transformers/examples/flax/language-modeling/目錄下讀者學(xué)完后可以掌握訓(xùn)練 tokenizer → 創(chuàng)建模型配置 → 啟動預(yù)訓(xùn)練 → 評估與推送 Hub的完整閉環(huán)并理解 JAX 純函數(shù)式編程、jax.pmap數(shù)據(jù)并行與 Optax 優(yōu)化器在語言模型訓(xùn)練中的實(shí)際落地方式。一、文檔定位Flax 語言模型預(yù)訓(xùn)練示例在倉庫中的角色該文檔隸屬于當(dāng)前倉庫benchmark/third_party/transformers/這份第三方 Transformers 代碼快照中。它與倉庫主項(xiàng)目 FlexGen面向單 GPU 吞吐優(yōu)先場景的大語言模型推理系統(tǒng)屬于參照與對比的關(guān)系FlexGen 專注于推理側(cè)的吞吐優(yōu)化而這些 Flax 示例則展示了完整的上游預(yù)訓(xùn)練鏈路。閱讀本指南時(shí)你既可以將這些腳本當(dāng)作獨(dú)立的語言模型訓(xùn)練工具使用也可以把它們視為理解大模型如何從無到有產(chǎn)生的背景知識。目錄中與文檔配套的腳本與文件包括run_mlm_flax.py掩碼語言建模MLM訓(xùn)練腳本支持 BERT/RoBERTa 等雙向模型run_clm_flax.py因果語言建模CLM訓(xùn)練腳本支持 GPT-2 等自回歸模型run_t5_mlm_flax.pyT5 式 span-masked 語言建模腳本run_bart_dlm_flax.pyBART 去噪語言建模腳本t5_tokenizer_model.py面向 T5 的 SentencePiece Unigram tokenizer 訓(xùn)練實(shí)現(xiàn)requirements.txt腳本運(yùn)行所需的最小依賴清單。二、為什么選擇 JAX/Flax 做預(yù)訓(xùn)練文檔開篇即點(diǎn)明了選型理由JAX/Flax 允許你trace 純函數(shù)并將其編譯成高效的、融合fused的加速器代碼可同時(shí)在 GPU 與 TPU 上運(yùn)行。用 JAX/Flax 編寫的模型是不可變的immutable參數(shù)以純函數(shù)式方式更新這使得模型并行model parallelism的實(shí)現(xiàn)簡單且高效——因?yàn)椴淮嬖诳勺儬顟B(tài)跨設(shè)備同步的問題。以 MLM 腳本為例訓(xùn)練主循環(huán)中可以看到這套范式在代碼層面的體現(xiàn)見run_mlm_flax.pyjax.value_and_grad(loss_fn, has_auxTrue)一次性返回?fù)p失與梯度jax.pmap(train_step, batch, donate_argnums(0,))將訓(xùn)練步編譯為跨設(shè)備并行的 SPMD 程序jax.lax.psum在 batch 維度上做跨設(shè)備歸約得到全局 loss 與梯度jax_utils.replicate(state)將訓(xùn)練狀態(tài)復(fù)制到每個(gè)設(shè)備。對訓(xùn)練吞吐有極致要求的團(tuán)隊(duì)例如 TPU pod 用戶這套 pipeline 能最大化加速器利用率這正是文檔選擇在單臺 TPUv3-8 上展示完整預(yù)訓(xùn)練流程的原因。三、環(huán)境準(zhǔn)備與依賴安裝所有示例腳本的依賴統(tǒng)一聲明在requirements.txt中最小版本如下datasets 1.1.3 jax0.2.8 jaxlib0.1.59 flax0.3.5 optax0.0.9此外還需要transformers、tokenizers以及用于訓(xùn)練日志可視化的tensorboard腳本通過is_tensorboard_available()檢測僅在主進(jìn)程寫入 SummaryWriter。若需要使用--push_to_hub還需通過huggingface-cli login完成 Hub 認(rèn)證。四、通用三步工作流任何模型都要先搭好模型倉庫無論訓(xùn)練哪種模型文檔都遵循同一套三步流程先在本地搭好一個(gè)模型倉庫目錄訓(xùn)練 tokenizer在目標(biāo)語料上訓(xùn)練分詞器并保存為tokenizer.json創(chuàng)建配置從 Hub 上現(xiàn)有 checkpoint 加載配置如roberta-base的RobertaConfig覆蓋vocab_size后保存到本地目錄訓(xùn)練模型運(yùn)行對應(yīng)的run_*_flax.py腳本訓(xùn)練過程中自動把日志與權(quán)重推送回該倉庫配合--push_to_hub。下面按文檔順序逐一展開四種任務(wù)的完整實(shí)操。五、掩碼語言建模MLM預(yù)訓(xùn)練 RoBERTa-base掩碼語言建模目標(biāo)源自 BERT 論文《BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding》。文檔示例在挪威語 OSCAR 語料上從零預(yù)訓(xùn)練roberta-base硬件環(huán)境為單臺 TPUv3-88 核。5.1 準(zhǔn)備目錄與訓(xùn)練 tokenizermkdir ./norwegian-roberta-base使用tokenizers庫的ByteLevelBPETokenizer與 Hugging Face 官方博客How to train a new language model from scratch一脈相承在 OSCAR 挪威語全量數(shù)據(jù)上訓(xùn)練from datasets import load_dataset from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer # load dataset dataset load_dataset(oscar, unshuffled_deduplicated_no, splittrain) # Instantiate tokenizer tokenizer ByteLevelBPETokenizer() def batch_iterator(batch_size1000): for i in range(0, len(dataset), batch_size): yield dataset[i: i batch_size][text] # Customized training tokenizer.train_from_iterator(batch_iterator(), vocab_size50265, min_frequency2, special_tokens[ s, pad, /s, unk, mask, ]) # Save files to disk tokenizer.save(./norwegian-roberta-base/tokenizer.json)要點(diǎn)說明vocab_size50265與 RoBERTa 原始詞表大小一致其中 5 個(gè)位置預(yù)留給特殊 tokenmin_frequency2過濾出現(xiàn)頻次過低的 tokenByteLevel BPE 天然兼容大小寫與多語言文本無需顯式 lowercasing文檔提示該步驟耗時(shí)約 10 分鐘視硬件而定。5.2 創(chuàng)建模型配置從 Hub 上的roberta-base加載配置模板僅覆蓋詞表大小from transformers import RobertaConfig config RobertaConfig.from_pretrained(roberta-base, vocab_size50265) config.save_pretrained(./norwegian-roberta-base)至此模型倉庫搭建完畢。訓(xùn)練期間腳本會自動把訓(xùn)練日志與模型權(quán)重推送到倉庫對應(yīng)run_mlm_flax.py中Repository的創(chuàng)建與repo.push_to_hub調(diào)用。5.3 啟動預(yù)訓(xùn)練python run_mlm_flax.py \ --output_dir./norwegian-roberta-base \ --model_typeroberta \ --config_name./norwegian-roberta-base \ --tokenizer_name./norwegian-roberta-base \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length128 \ --weight_decay0.01 \ --per_device_train_batch_size128 \ --per_device_eval_batch_size128 \ --learning_rate3e-4 \ --warmup_steps1000 \ --overwrite_output_dir \ --num_train_epochs18 \ --adam_beta10.9 \ --adam_beta20.98 \ --logging_steps500 \ --save_steps2500 \ --eval_steps2500 \ --push_to_hub文檔給出的收斂參考在單臺 TPUv3-8 上訓(xùn)練 18 個(gè) epoch 后loss 與準(zhǔn)確率分別收斂到1.78 與 0.64總耗時(shí)約 18 小時(shí)以內(nèi)。5.4 源碼層面的掩碼策略與訓(xùn)練機(jī)制run_mlm_flax.py中FlaxDataCollatorForLanguageModeling.mask_tokensbenchmark/third_party/transformers/examples/flax/language-modeling/run_mlm_flax.py實(shí)現(xiàn)了標(biāo)準(zhǔn)的 80/10/10 掩碼策略以mlm_probability默認(rèn) 0.15概率隨機(jī)選出待掩碼 token其中 80% 替換為[MASK]tokenizer.mask_token10% 替換為詞表中的隨機(jī)詞剩余 10% 保持不變未被選中的 token 對應(yīng) label 置為-100只對掩碼 token 計(jì)算 loss特殊 tokenspecial_tokens_mask標(biāo)記永遠(yuǎn)不會被掩碼。訓(xùn)練循環(huán)的其他關(guān)鍵機(jī)制包括學(xué)習(xí)率調(diào)度optax.linear_schedule構(gòu)建線性 warmup 到learning_rate再線性衰減到 0的雙段調(diào)度邊界在warmup_steps處拼接optax.join_schedules優(yōu)化器默認(rèn)optax.adamwAdamWdecay_mask_fn通過traverse_util.flatten_dict找出所有 LayerNorm 參數(shù)與 bias構(gòu)造 mask 使權(quán)重衰減不作用于 bias 與 LayerNorm 參數(shù)也可用--adafactor切換為optax.adafactor數(shù)據(jù)并行train_batch_size per_device_train_batch_size * jax.device_count()shard切分 batch 到各設(shè)備pad_shard_unpad處理評估時(shí)的不整除檢查點(diǎn)與推送每save_steps步主進(jìn)程執(zhí)行model.save_pretrainedtokenizer.save_pretrained若開啟--push_to_hub則異步推送blockingFalse評估評估時(shí)額外計(jì)算 accuracy并在訓(xùn)練結(jié)束后把eval_results.json寫入輸出目錄包含perplexity exp(loss)。六、因果語言建模CLM預(yù)訓(xùn)練 GPT-2第二部分演示自回歸因果語言模型的訓(xùn)練在挪威語 OSCAR 語料上從零預(yù)訓(xùn)練隨機(jī)初始化的 124M 參數(shù)gpt2同樣運(yùn)行在單臺 TPUv3-8 上。6.1 訓(xùn)練 tokenizermkdir ./norwegian-gpt2from datasets import load_dataset from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer # load dataset dataset load_dataset(oscar, unshuffled_deduplicated_no, splittrain) # Instantiate tokenizer tokenizer ByteLevelBPETokenizer() def batch_iterator(batch_size1000): for i in range(0, len(dataset), batch_size): yield dataset[i: i batch_size][text] # Customized training tokenizer.train_from_iterator(batch_iterator(), vocab_size50257, min_frequency2, special_tokens[ s, pad, /s, unk, mask, ]) # Save files to disk tokenizer.save(./norwegian-gpt2/tokenizer.json)與 MLM 不同之處在于vocab_size50257GPT-2 原始詞表大小。6.2 創(chuàng)建配置from transformers import GPT2Config config GPT2Config.from_pretrained(gpt2, resid_pdrop0.0, embd_pdrop0.0, attn_pdrop0.0, vocab_size50257) config.save_pretrained(./norwegian-gpt2)注意這里顯式將三類 dropoutresid_pdrop、embd_pdrop、attn_pdrop置為 0——預(yù)訓(xùn)練從零開始時(shí)通常不啟用 dropout避免隨機(jī)失活干擾梯度信號。6.3 啟動預(yù)訓(xùn)練python run_clm_flax.py \ --output_dir./norwegian-gpt2 \ --model_typegpt2 \ --config_name./norwegian-gpt2 \ --tokenizer_name./norwegian-gpt2 \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --do_train --do_eval \ --block_size512 \ --per_device_train_batch_size64 \ --per_device_eval_batch_size64 \ --learning_rate5e-3 --warmup_steps1000 \ --adam_beta10.9 --adam_beta20.98 --weight_decay0.01 \ --overwrite_output_dir \ --num_train_epochs20 \ --logging_steps500 \ --save_steps2500 \ --eval_steps2500 \ --push_to_hub與 MLM 命令的關(guān)鍵差異--block_size512CLM 無需掩碼語料按固定窗口切分為連續(xù) token 塊run_clm_flax.py中的group_texts會把全部文本拼接后切成block_size的塊學(xué)習(xí)率顯著更高5e-3因?yàn)樽曰貧w目標(biāo)的訓(xùn)練信號更密集顯式給出--do_train --do_eval。文檔給出的收斂參考20 個(gè) epoch 后 loss 與困惑度分別收斂到3.24 與 25.72耗時(shí)約 21 小時(shí)以內(nèi)。評估階段的困惑度即exp(eval_loss)與run_mlm_flax.py末尾的eval_results.json計(jì)算邏輯一致。七、T5 式 span-masked 語言建模預(yù)訓(xùn)練 T5 v1.1-base第三部分演示以 T5 論文《Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》提出的 span-masked 目標(biāo)預(yù)訓(xùn)練google/t5-v1_1-base。T5 是 encoder-decoder 結(jié)構(gòu)掩碼對象是連續(xù)的 token 片段span模型需要根據(jù)上下文重建被掩碼的整段內(nèi)容。7.1 訓(xùn)練 SentencePiece Unigram tokenizerT5 家族使用 SentencePiece 風(fēng)格的 tokenizer與 BPE 不同。倉庫提供t5_tokenizer_model.py中的SentencePieceUnigramTokenizer其實(shí)現(xiàn)細(xì)節(jié)benchmark/third_party/transformers/examples/flax/language-modeling/t5_tokenizer_model.py包括采用Unigram 語言模型算法tokenizers.models.Unigram歸一化流水線Nmt→NFKC→ 將連續(xù)多個(gè)空格壓縮為單個(gè) →Lowercase預(yù)分詞流水線Metaspace以▁表示空格并默認(rèn)在句首加前綴空格→Digits(individual_digitsTrue)數(shù)字逐位拆分→Punctuation標(biāo)點(diǎn)獨(dú)立成 token解碼器與后處理器Metaspace解碼單序列自動追加/sTemplateProcessing特殊 token 約定pad0、/s1、unk2訓(xùn)練結(jié)束后通過add_unk_id把 unk id 寫回模型配置。訓(xùn)練腳本如下注意目錄創(chuàng)建方式為cd進(jìn)入cd ./norwegian-t5-baseimport datasets from t5_tokenizer_model import SentencePieceUnigramTokenizer vocab_size 32_000 input_sentence_size None # Initialize a dataset dataset datasets.load_dataset(oscar, nameunshuffled_deduplicated_no, splittrain) tokenizer SentencePieceUnigramTokenizer(unk_tokenunk, eos_token/s, pad_tokenpad) # Build an iterator over this dataset def batch_iterator(input_sentence_sizeNone): if input_sentence_size is None: input_sentence_size len(dataset) batch_length 100 for i in range(0, input_sentence_size, batch_length): yield dataset[i: i batch_length][text] # Train tokenizer tokenizer.train_from_iterator( iteratorbatch_iterator(input_sentence_sizeinput_sentence_size), vocab_sizevocab_size, show_progressTrue, ) # Save files to disk tokenizer.save(./norwegian-t5-base/tokenizer.json)T5 詞表大小為 32,000。由于 Unigram 訓(xùn)練需要多輪 EM 迭代文檔提示該步驟最長可達(dá) 120 分鐘是四種模型中耗時(shí)最長的。7.2 創(chuàng)建配置from transformers import T5Config config T5Config.from_pretrained(google/t5-v1_1-base, vocab_sizetokenizer.get_vocab_size()) config.save_pretrained(./norwegian-t5-base)這里vocab_size直接取 tokenizer 實(shí)際訓(xùn)練出的詞表大小保證詞嵌入維度與 tokenizer 嚴(yán)格對齊。7.3 啟動預(yù)訓(xùn)練python run_t5_mlm_flax.py \ --output_dir./norwegian-t5-base \ --model_typet5 \ --config_name./norwegian-t5-base \ --tokenizer_name./norwegian-t5-base \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length512 \ --per_device_train_batch_size32 \ --per_device_eval_batch_size32 \ --adafactor \ --learning_rate0.005 \ --weight_decay0.001 \ --warmup_steps2000 \ --overwrite_output_dir \ --logging_steps500 \ --save_steps10000 \ --eval_steps2500 \ --push_to_hub關(guān)鍵差異--adafactorT5 訓(xùn)練慣例上使用 Adafactor 而非 AdamW這是文檔明確給出的配置也是run_t5_mlm_flax.py中TrainingArguments.adafactor開關(guān)的典型用法max_seq_length512、batch size 32、save_steps10000間隔更大因?yàn)?T5 單步更耗時(shí)。文檔給出的收斂參考3 個(gè) epoch 后 loss 與準(zhǔn)確率收斂到2.36 與 57.0耗時(shí)約 4.5 小時(shí)。訓(xùn)練統(tǒng)計(jì)可直接在 Hub 模型倉庫的 TensorBoard 頁查看。八、BART 去噪語言建模DLM預(yù)訓(xùn)練 BART-base第四部分演示以 BART 論文《BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension》提出的去噪目標(biāo)預(yù)訓(xùn)練facebook/bart-base。BART 同樣是 encoder-decoder 架構(gòu)但與 T5 的 span-masking 不同它的輸入會疊加多種噪聲如 token 刪除、文本旋轉(zhuǎn)、span 掩碼、句序打亂目標(biāo)是讓 decoder 重建完整原文。8.1 訓(xùn)練 tokenizer 與創(chuàng)建配置mkdir ./norwegian-bart-basetokenizer 訓(xùn)練與 RoBERTa 完全一致ByteLevel BPEvocab_size50265代碼可復(fù)用 5.1 節(jié)示例。配置創(chuàng)建from transformers import BartConfig config BartConfig.from_pretrained(facebook/bart-base, vocab_size50265) config.save_pretrained(./norwegian-bart-base)8.2 啟動預(yù)訓(xùn)練python run_bart_dlm_flax.py \ --output_dir./norwegian-bart-base \ --config_name./norwegian-bart-base \ --tokenizer_name./norwegian-bart-base \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length1024 \ --per_device_train_batch_size32 \ --per_device_eval_batch_size32 \ --learning_rate1e-4 \ --warmup_steps2000 \ --overwrite_output_dir \ --logging_steps500 \ --save_steps2000 \ --eval_steps2000 \ --push_to_hub注意 BART 的max_seq_length1024是四種模型中最長的上下文窗口。文檔給出的收斂參考3 個(gè) epoch 后 loss 與準(zhǔn)確率收斂到1.36 與 0.77耗時(shí) 6 小時(shí)以內(nèi)。九、訓(xùn)練參數(shù)全景表來自腳本源碼run_mlm_flax.py、run_clm_flax.py、run_t5_mlm_flax.py中三組 dataclass 定義了全部命令行參數(shù)。下表匯總默認(rèn)值與含義除特別說明外三個(gè)腳本一致TrainingArguments參數(shù)默認(rèn)值說明--output_dir必填預(yù)測結(jié)果與 checkpoint 輸出目錄--overwrite_output_dirFalse輸出目錄非空時(shí)覆蓋指向已有 checkpoint 目錄可續(xù)訓(xùn)--do_train/--do_evalFalse是否執(zhí)行訓(xùn)練 / 評估--per_device_train_batch_size8每個(gè) GPU/TPU 核/CPU 的訓(xùn)練 batch--per_device_eval_batch_size8每個(gè)設(shè)備的評估 batch--learning_rate5e-5初始學(xué)習(xí)率--weight_decay0.0AdamW 權(quán)重衰減系數(shù)--adam_beta10.9Adam 一階矩指數(shù)衰減--adam_beta20.999Adam 二階矩指數(shù)衰減--adam_epsilon1e-8Adam epsilon--adafactorFalse用 Adafactor 替代 AdamW--num_train_epochs3.0訓(xùn)練總 epoch 數(shù)--warmup_steps0線性 warmup 步數(shù)--logging_steps500每 X 步記錄一次日志--save_steps500每 X 步保存一次 checkpoint--eval_stepsNone每 X 步執(zhí)行一次評估--seed42全局隨機(jī)種子--push_to_hubFalse訓(xùn)練后是否上傳模型到 Hub--hub_model_id/--hub_tokenNoneHub 倉庫名 / 認(rèn)證 token--gradient_checkpointingFalse梯度檢查點(diǎn)省顯存換速度僅 MLM 腳本支持ModelArguments參數(shù)默認(rèn)值說明--model_name_or_pathNone初始化權(quán)重的 checkpoint不傳則從零訓(xùn)練--model_typeNone從零訓(xùn)練時(shí)指定模型類型從FLAX_MODEL_FOR_*_LM_MAPPING鍵中取--config_name/--tokenizer_nameNone與模型不同的 config / tokenizer 名稱或路徑--cache_dirNone預(yù)訓(xùn)練模型緩存目錄--use_fast_tokenizerTrue使用 fast tokenizer基于 tokenizers 庫--dtypefloat32權(quán)重初始化與訓(xùn)練精度float32/float16/bfloat16--use_auth_tokenFalse訪問私有模型所需的 Hub tokenDataTrainingArguments以 MLM 腳本為例參數(shù)默認(rèn)值說明--dataset_name/--dataset_config_nameNone通過 datasets 庫加載的數(shù)據(jù)集名與子配置名--train_file/--validation_fileNone本地 csv/json/txt 訓(xùn)練與驗(yàn)證文件--validation_split_percentage5無驗(yàn)證集時(shí)從訓(xùn)練集切出的驗(yàn)證比例--max_seq_lengthNone最大序列長度超過則截?cái)嗄J(rèn)取模型上限--preprocessing_num_workersNone預(yù)處理并行進(jìn)程數(shù)--mlm_probability0.15MLM 掩碼比例--pad_to_max_lengthFalse是否全部 padding 到max_seq_length--line_by_lineFalse是否按行切分序列否則拼接全部文本再切塊實(shí)戰(zhàn)提示腳本同時(shí)支持把參數(shù)寫成 JSON 文件傳入python run_mlm_flax.py args.json源碼通過HfArgumentParser.parse_json_file支持便于在集群上固化實(shí)驗(yàn)配置。十、運(yùn)行時(shí)評估Flax vs PyTorch/XLA vs PyTorch文檔對同一個(gè) MLM 任務(wù)在三種后端/硬件組合下做了耗時(shí)對比用于說明 Flax 在 TPU 上的編譯與并行優(yōu)勢。所有實(shí)驗(yàn)均運(yùn)行在 Google Cloud Platform 上任務(wù)TPU v3-8FlaxTPU v3-8PyTorch/XLA8 GPUPyTorchMLM15h32m23h46m44h14m說明GPU 實(shí)驗(yàn)除 JAX 變換外未做額外優(yōu)化且以全精度 fp32 運(yùn)行TPU v3-8 指 4 塊芯片共 8 個(gè) TPU 核8 GPU 指 8 塊獨(dú)立 GPU 芯片。10.1 用 PyTorch/XLA 在 TPUv3-8 上復(fù)現(xiàn)同一實(shí)驗(yàn)前提是已按 5.1/5.2 節(jié)在norwegian-roberta-base目錄備好 tokenizer 與配置。在示例目錄下建立符號鏈接把 PyTorch 版run_mlm.py與xla_spawn.py引入當(dāng)前工作目錄ln -s ~/transformers/examples/pytorch/language-modeling/run_mlm.py ./ ln -s ~/transformers/examples/pytorch/xla_spawn.py ./設(shè)置環(huán)境變量export XRT_TPU_CONFIGlocalservice;0;localhost:51011 unset LD_PRELOAD export NUM_TPUS8 export TOKENIZERS_PARALLELISM0 export MODEL_DIR./norwegian-roberta-base mkdir -p ${MODEL_DIR}啟動訓(xùn)練python3 xla_spawn.py --num_cores ${NUM_TPUS} run_mlm.py --output_dir./runs \ --model_typeroberta \ --config_name${MODEL_DIR} \ --tokenizer_name${MODEL_DIR} \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length128 \ --weight_decay0.01 \ --per_device_train_batch_size128 \ --per_device_eval_batch_size128 \ --learning_rate3e-4 \ --warmup_steps1000 \ --overwrite_output_dir \ --num_train_epochs18 \ --adam_beta10.9 \ --adam_beta20.98 \ --do_train \ --do_eval \ --logging_steps500 \ --evaluation_strategyepoch \ --report_totensorboard \ --save_strategyno10.2 用 PyTorch 在 8 張 V100 GPU 上復(fù)現(xiàn)同一實(shí)驗(yàn)由于單張 V100 顯存限制最大 batch 只能到 32Flax 版為 128必須引入gradient_accumulation彌補(bǔ)全局 batchln -s ~/transformers/examples/pytorch/language-modeling/run_mlm.py ./export NUM_GPUS8 export TOKENIZERS_PARALLELISM0 export MODEL_DIR./norwegian-roberta-base mkdir -p ${MODEL_DIR}python3 -m torch.distributed.launch --nproc_per_node ${NUM_GPUS} run_mlm.py \ --output_dir${MODEL_DIR} \ --model_typeroberta \ --config_name${MODEL_DIR} \ --tokenizer_name${MODEL_DIR} \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length128 \ --weight_decay0.01 \ --per_device_train_batch_size32 \ --per_device_eval_batch_size32 \ --gradient_accumulation4 \ --learning_rate3e-4 \ --warmup_steps1000 \ --overwrite_output_dir \ --num_train_epochs18 \ --adam_beta10.9 \ --adam_beta20.98 \ --do_train \ --do_eval \ --logging_steps500 \ --evaluation_strategysteps \ --report_totensorboard \ --save_strategynoper_device_train_batch_size32乘以gradient_accumulation4再乘以 8 卡全局有效 batch 與 Flax 版128×8等價(jià)保證對比公平。十一、從預(yù)訓(xùn)練到推理與 FlexGen 的關(guān)系完成上述任一預(yù)訓(xùn)練流程后得到的 checkpoint 即可被加載用于推理/微調(diào)。這與倉庫主項(xiàng)目 FlexGen 構(gòu)成完整的上下游鏈條FlexGen 專注于讓大語言模型在單 GPU 上以吞吐優(yōu)先的方式高效推理見倉庫根目錄 README.md 與flexgen/下的實(shí)現(xiàn)而本文所述 Flax 示例則負(fù)責(zé)模型從零到有的預(yù)訓(xùn)練環(huán)節(jié)。對于想在生產(chǎn)環(huán)境落地大模型的讀者可以先用本文流程訓(xùn)練領(lǐng)域語料模型再借助 FlexGen 的思路做推理側(cè)的吞吐優(yōu)化兩種能力互為補(bǔ)充覆蓋了從訓(xùn)練到服務(wù)的完整生命周期。十二、常見問題與調(diào)參建議收斂不佳時(shí)優(yōu)先檢查什么確認(rèn)vocab_size與 tokenizer 實(shí)際詞表一致T5 場景務(wù)必用tokenizer.get_vocab_size()確認(rèn)max_seq_length不超過模型支持上限源碼中max_seq_length min(data_args.max_seq_length, tokenizer.model_max_length)顯存不足降低per_device_*_batch_size或開啟--gradient_checkpointing僅run_mlm_flax.py提供該開關(guān)PyTorch 對比實(shí)驗(yàn)則用gradient_accumulation補(bǔ)償多機(jī)多卡Flax 腳本依賴jax.device_count()自動感知設(shè)備數(shù)訓(xùn)練 batch 會按設(shè)備數(shù)自動放大無需手工配置分布式啟動器自定義語料不需要 Hub 數(shù)據(jù)集時(shí)直接用--train_file/--validation_file傳入 csv/json/txt 文件即可腳本會自動按text列或第一列讀取復(fù)現(xiàn)性--seed默認(rèn) 42腳本在模型初始化與數(shù)據(jù)洗牌處均注入該種子。本文所有腳本、配置與依賴均可在當(dāng)前倉庫benchmark/third_party/transformers/examples/flax/language-modeling/目錄下找到并直接運(yùn)行建議結(jié)合源碼閱讀本文以獲得對 JAX/Flax 訓(xùn)練管線最完整的理解。贊分享推理引擎大模型【免費(fèi)下載鏈接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.項(xiàng)目地址https://gitcode.com/gh_mirrors/fl/FlexGen點(diǎn)擊查看免費(fèi)下載相關(guān)推薦使用 Ray Train JaxTrainer 從零分布式訓(xùn)練 GPT-2JAX/Flax 規(guī)?；?xùn)練實(shí)戰(zhàn)指南使用 Ray Train JaxTrainer 從零分布式訓(xùn)練 GPT 2JAX/Flax 規(guī)?；?xùn)練實(shí)戰(zhàn)指南 預(yù)計(jì)閱讀與實(shí)操時(shí)間 約 15 分鐘 本指南以人工智能分布式訓(xùn)練強(qiáng)化學(xué)習(xí)任務(wù)調(diào)度模型推理服務(wù)后端使用 SpeechBrain 訓(xùn)練 Switchboard 語言模型Transformer LM 從零訓(xùn)練與預(yù)訓(xùn)練微調(diào)實(shí)戰(zhàn)指南使用 SpeechBrain 訓(xùn)練 Switchboard 語言模型Transformer LM 從零訓(xùn)練與預(yù)訓(xùn)練微調(diào)實(shí)戰(zhàn)指南 本篇技術(shù)指南以 SpeechB人工智能深度學(xué)習(xí)語音音頻NLP預(yù)訓(xùn)練Fairseq 多語言 RoBERTa 預(yù)訓(xùn)練實(shí)戰(zhàn)指南從多語語料預(yù)處理到 masked LM 訓(xùn)練Fairseq 多語言 RoBERTa 預(yù)訓(xùn)練實(shí)戰(zhàn)指南從多語語料預(yù)處理到 masked LM 訓(xùn)練 多語言 RoBERTaMultilingual RoBE人工智能深度學(xué)習(xí)預(yù)訓(xùn)練NLP語音上一篇FastGPT 工作流變量替換 CPU 阻塞優(yōu)化從變量驅(qū)動到模板驅(qū)動的調(diào)度熱路徑重構(gòu)下一篇LeetCode 179 Largest Number 的 Go 題解自定義比較器快排與 ab、ba 拼接比較創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 FlexGen 倉庫視角下的 Transformers Flax 語言模型預(yù)訓(xùn)練實(shí)戰(zhàn)：從零訓(xùn)練 RoBERTa、GPT-2、T5 與 BART
推理引擎大模型【免費(fèi)下載鏈接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.項(xiàng)目地址https://gitcode.com/gh_mirrors/fl/FlexGen點(diǎn)擊查看免費(fèi)下載本篇技術(shù)指南以benchmark/third_party/transformers/examples/flax/language-modeling/README.md為核心文檔系統(tǒng)講解如何利用 JAX/Flax 后端從零預(yù)訓(xùn)練四種主流語言模型——雙向的 RoBERTaMLM、自回歸的 GPT-2CLM、span-masked 的 T5 以及去噪自編碼的 BART。文中的示例腳本與配套代碼均位于當(dāng)前倉庫的benchmark/third_party/transformers/examples/flax/language-modeling/目錄下讀者學(xué)完后可以掌握訓(xùn)練 tokenizer → 創(chuàng)建模型配置 → 啟動預(yù)訓(xùn)練 → 評估與推送 Hub的完整閉環(huán)并理解 JAX 純函數(shù)式編程、jax.pmap數(shù)據(jù)并行與 Optax 優(yōu)化器在語言模型訓(xùn)練中的實(shí)際落地方式。一、文檔定位Flax 語言模型預(yù)訓(xùn)練示例在倉庫中的角色該文檔隸屬于當(dāng)前倉庫benchmark/third_party/transformers/這份第三方 Transformers 代碼快照中。它與倉庫主項(xiàng)目 FlexGen面向單 GPU 吞吐優(yōu)先場景的大語言模型推理系統(tǒng)屬于參照與對比的關(guān)系FlexGen 專注于推理側(cè)的吞吐優(yōu)化而這些 Flax 示例則展示了完整的上游預(yù)訓(xùn)練鏈路。閱讀本指南時(shí)你既可以將這些腳本當(dāng)作獨(dú)立的語言模型訓(xùn)練工具使用也可以把它們視為理解大模型如何從無到有產(chǎn)生的背景知識。目錄中與文檔配套的腳本與文件包括run_mlm_flax.py掩碼語言建模MLM訓(xùn)練腳本支持 BERT/RoBERTa 等雙向模型run_clm_flax.py因果語言建模CLM訓(xùn)練腳本支持 GPT-2 等自回歸模型run_t5_mlm_flax.pyT5 式 span-masked 語言建模腳本run_bart_dlm_flax.pyBART 去噪語言建模腳本t5_tokenizer_model.py面向 T5 的 SentencePiece Unigram tokenizer 訓(xùn)練實(shí)現(xiàn)requirements.txt腳本運(yùn)行所需的最小依賴清單。二、為什么選擇 JAX/Flax 做預(yù)訓(xùn)練文檔開篇即點(diǎn)明了選型理由JAX/Flax 允許你trace 純函數(shù)并將其編譯成高效的、融合fused的加速器代碼可同時(shí)在 GPU 與 TPU 上運(yùn)行。用 JAX/Flax 編寫的模型是不可變的immutable參數(shù)以純函數(shù)式方式更新這使得模型并行model parallelism的實(shí)現(xiàn)簡單且高效——因?yàn)椴淮嬖诳勺儬顟B(tài)跨設(shè)備同步的問題。以 MLM 腳本為例訓(xùn)練主循環(huán)中可以看到這套范式在代碼層面的體現(xiàn)見run_mlm_flax.pyjax.value_and_grad(loss_fn, has_auxTrue)一次性返回?fù)p失與梯度jax.pmap(train_step, batch, donate_argnums(0,))將訓(xùn)練步編譯為跨設(shè)備并行的 SPMD 程序jax.lax.psum在 batch 維度上做跨設(shè)備歸約得到全局 loss 與梯度jax_utils.replicate(state)將訓(xùn)練狀態(tài)復(fù)制到每個(gè)設(shè)備。對訓(xùn)練吞吐有極致要求的團(tuán)隊(duì)例如 TPU pod 用戶這套 pipeline 能最大化加速器利用率這正是文檔選擇在單臺 TPUv3-8 上展示完整預(yù)訓(xùn)練流程的原因。三、環(huán)境準(zhǔn)備與依賴安裝所有示例腳本的依賴統(tǒng)一聲明在requirements.txt中最小版本如下datasets 1.1.3 jax0.2.8 jaxlib0.1.59 flax0.3.5 optax0.0.9此外還需要transformers、tokenizers以及用于訓(xùn)練日志可視化的tensorboard腳本通過is_tensorboard_available()檢測僅在主進(jìn)程寫入 SummaryWriter。若需要使用--push_to_hub還需通過huggingface-cli login完成 Hub 認(rèn)證。四、通用三步工作流任何模型都要先搭好模型倉庫無論訓(xùn)練哪種模型文檔都遵循同一套三步流程先在本地搭好一個(gè)模型倉庫目錄訓(xùn)練 tokenizer在目標(biāo)語料上訓(xùn)練分詞器并保存為tokenizer.json創(chuàng)建配置從 Hub 上現(xiàn)有 checkpoint 加載配置如roberta-base的RobertaConfig覆蓋vocab_size后保存到本地目錄訓(xùn)練模型運(yùn)行對應(yīng)的run_*_flax.py腳本訓(xùn)練過程中自動把日志與權(quán)重推送回該倉庫配合--push_to_hub。下面按文檔順序逐一展開四種任務(wù)的完整實(shí)操。五、掩碼語言建模MLM預(yù)訓(xùn)練 RoBERTa-base掩碼語言建模目標(biāo)源自 BERT 論文《BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding》。文檔示例在挪威語 OSCAR 語料上從零預(yù)訓(xùn)練roberta-base硬件環(huán)境為單臺 TPUv3-88 核。5.1 準(zhǔn)備目錄與訓(xùn)練 tokenizermkdir ./norwegian-roberta-base使用tokenizers庫的ByteLevelBPETokenizer與 Hugging Face 官方博客How to train a new language model from scratch一脈相承在 OSCAR 挪威語全量數(shù)據(jù)上訓(xùn)練from datasets import load_dataset from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer # load dataset dataset load_dataset(oscar, unshuffled_deduplicated_no, splittrain) # Instantiate tokenizer tokenizer ByteLevelBPETokenizer() def batch_iterator(batch_size1000): for i in range(0, len(dataset), batch_size): yield dataset[i: i batch_size][text] # Customized training tokenizer.train_from_iterator(batch_iterator(), vocab_size50265, min_frequency2, special_tokens[ s, pad, /s, unk, mask, ]) # Save files to disk tokenizer.save(./norwegian-roberta-base/tokenizer.json)要點(diǎn)說明vocab_size50265與 RoBERTa 原始詞表大小一致其中 5 個(gè)位置預(yù)留給特殊 tokenmin_frequency2過濾出現(xiàn)頻次過低的 tokenByteLevel BPE 天然兼容大小寫與多語言文本無需顯式 lowercasing文檔提示該步驟耗時(shí)約 10 分鐘視硬件而定。5.2 創(chuàng)建模型配置從 Hub 上的roberta-base加載配置模板僅覆蓋詞表大小from transformers import RobertaConfig config RobertaConfig.from_pretrained(roberta-base, vocab_size50265) config.save_pretrained(./norwegian-roberta-base)至此模型倉庫搭建完畢。訓(xùn)練期間腳本會自動把訓(xùn)練日志與模型權(quán)重推送到倉庫對應(yīng)run_mlm_flax.py中Repository的創(chuàng)建與repo.push_to_hub調(diào)用。5.3 啟動預(yù)訓(xùn)練python run_mlm_flax.py \ --output_dir./norwegian-roberta-base \ --model_typeroberta \ --config_name./norwegian-roberta-base \ --tokenizer_name./norwegian-roberta-base \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length128 \ --weight_decay0.01 \ --per_device_train_batch_size128 \ --per_device_eval_batch_size128 \ --learning_rate3e-4 \ --warmup_steps1000 \ --overwrite_output_dir \ --num_train_epochs18 \ --adam_beta10.9 \ --adam_beta20.98 \ --logging_steps500 \ --save_steps2500 \ --eval_steps2500 \ --push_to_hub文檔給出的收斂參考在單臺 TPUv3-8 上訓(xùn)練 18 個(gè) epoch 后loss 與準(zhǔn)確率分別收斂到1.78 與 0.64總耗時(shí)約 18 小時(shí)以內(nèi)。5.4 源碼層面的掩碼策略與訓(xùn)練機(jī)制run_mlm_flax.py中FlaxDataCollatorForLanguageModeling.mask_tokensbenchmark/third_party/transformers/examples/flax/language-modeling/run_mlm_flax.py實(shí)現(xiàn)了標(biāo)準(zhǔn)的 80/10/10 掩碼策略以mlm_probability默認(rèn) 0.15概率隨機(jī)選出待掩碼 token其中 80% 替換為[MASK]tokenizer.mask_token10% 替換為詞表中的隨機(jī)詞剩余 10% 保持不變未被選中的 token 對應(yīng) label 置為-100只對掩碼 token 計(jì)算 loss特殊 tokenspecial_tokens_mask標(biāo)記永遠(yuǎn)不會被掩碼。訓(xùn)練循環(huán)的其他關(guān)鍵機(jī)制包括學(xué)習(xí)率調(diào)度optax.linear_schedule構(gòu)建線性 warmup 到learning_rate再線性衰減到 0的雙段調(diào)度邊界在warmup_steps處拼接optax.join_schedules優(yōu)化器默認(rèn)optax.adamwAdamWdecay_mask_fn通過traverse_util.flatten_dict找出所有 LayerNorm 參數(shù)與 bias構(gòu)造 mask 使權(quán)重衰減不作用于 bias 與 LayerNorm 參數(shù)也可用--adafactor切換為optax.adafactor數(shù)據(jù)并行train_batch_size per_device_train_batch_size * jax.device_count()shard切分 batch 到各設(shè)備pad_shard_unpad處理評估時(shí)的不整除檢查點(diǎn)與推送每save_steps步主進(jìn)程執(zhí)行model.save_pretrainedtokenizer.save_pretrained若開啟--push_to_hub則異步推送blockingFalse評估評估時(shí)額外計(jì)算 accuracy并在訓(xùn)練結(jié)束后把eval_results.json寫入輸出目錄包含perplexity exp(loss)。六、因果語言建模CLM預(yù)訓(xùn)練 GPT-2第二部分演示自回歸因果語言模型的訓(xùn)練在挪威語 OSCAR 語料上從零預(yù)訓(xùn)練隨機(jī)初始化的 124M 參數(shù)gpt2同樣運(yùn)行在單臺 TPUv3-8 上。6.1 訓(xùn)練 tokenizermkdir ./norwegian-gpt2from datasets import load_dataset from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer # load dataset dataset load_dataset(oscar, unshuffled_deduplicated_no, splittrain) # Instantiate tokenizer tokenizer ByteLevelBPETokenizer() def batch_iterator(batch_size1000): for i in range(0, len(dataset), batch_size): yield dataset[i: i batch_size][text] # Customized training tokenizer.train_from_iterator(batch_iterator(), vocab_size50257, min_frequency2, special_tokens[ s, pad, /s, unk, mask, ]) # Save files to disk tokenizer.save(./norwegian-gpt2/tokenizer.json)與 MLM 不同之處在于vocab_size50257GPT-2 原始詞表大小。6.2 創(chuàng)建配置from transformers import GPT2Config config GPT2Config.from_pretrained(gpt2, resid_pdrop0.0, embd_pdrop0.0, attn_pdrop0.0, vocab_size50257) config.save_pretrained(./norwegian-gpt2)注意這里顯式將三類 dropoutresid_pdrop、embd_pdrop、attn_pdrop置為 0——預(yù)訓(xùn)練從零開始時(shí)通常不啟用 dropout避免隨機(jī)失活干擾梯度信號。6.3 啟動預(yù)訓(xùn)練python run_clm_flax.py \ --output_dir./norwegian-gpt2 \ --model_typegpt2 \ --config_name./norwegian-gpt2 \ --tokenizer_name./norwegian-gpt2 \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --do_train --do_eval \ --block_size512 \ --per_device_train_batch_size64 \ --per_device_eval_batch_size64 \ --learning_rate5e-3 --warmup_steps1000 \ --adam_beta10.9 --adam_beta20.98 --weight_decay0.01 \ --overwrite_output_dir \ --num_train_epochs20 \ --logging_steps500 \ --save_steps2500 \ --eval_steps2500 \ --push_to_hub與 MLM 命令的關(guān)鍵差異--block_size512CLM 無需掩碼語料按固定窗口切分為連續(xù) token 塊run_clm_flax.py中的group_texts會把全部文本拼接后切成block_size的塊學(xué)習(xí)率顯著更高5e-3因?yàn)樽曰貧w目標(biāo)的訓(xùn)練信號更密集顯式給出--do_train --do_eval。文檔給出的收斂參考20 個(gè) epoch 后 loss 與困惑度分別收斂到3.24 與 25.72耗時(shí)約 21 小時(shí)以內(nèi)。評估階段的困惑度即exp(eval_loss)與run_mlm_flax.py末尾的eval_results.json計(jì)算邏輯一致。七、T5 式 span-masked 語言建模預(yù)訓(xùn)練 T5 v1.1-base第三部分演示以 T5 論文《Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》提出的 span-masked 目標(biāo)預(yù)訓(xùn)練google/t5-v1_1-base。T5 是 encoder-decoder 結(jié)構(gòu)掩碼對象是連續(xù)的 token 片段span模型需要根據(jù)上下文重建被掩碼的整段內(nèi)容。7.1 訓(xùn)練 SentencePiece Unigram tokenizerT5 家族使用 SentencePiece 風(fēng)格的 tokenizer與 BPE 不同。倉庫提供t5_tokenizer_model.py中的SentencePieceUnigramTokenizer其實(shí)現(xiàn)細(xì)節(jié)benchmark/third_party/transformers/examples/flax/language-modeling/t5_tokenizer_model.py包括采用Unigram 語言模型算法tokenizers.models.Unigram歸一化流水線Nmt→NFKC→ 將連續(xù)多個(gè)空格壓縮為單個(gè) →Lowercase預(yù)分詞流水線Metaspace以▁表示空格并默認(rèn)在句首加前綴空格→Digits(individual_digitsTrue)數(shù)字逐位拆分→Punctuation標(biāo)點(diǎn)獨(dú)立成 token解碼器與后處理器Metaspace解碼單序列自動追加/sTemplateProcessing特殊 token 約定pad0、/s1、unk2訓(xùn)練結(jié)束后通過add_unk_id把 unk id 寫回模型配置。訓(xùn)練腳本如下注意目錄創(chuàng)建方式為cd進(jìn)入cd ./norwegian-t5-baseimport datasets from t5_tokenizer_model import SentencePieceUnigramTokenizer vocab_size 32_000 input_sentence_size None # Initialize a dataset dataset datasets.load_dataset(oscar, nameunshuffled_deduplicated_no, splittrain) tokenizer SentencePieceUnigramTokenizer(unk_tokenunk, eos_token/s, pad_tokenpad) # Build an iterator over this dataset def batch_iterator(input_sentence_sizeNone): if input_sentence_size is None: input_sentence_size len(dataset) batch_length 100 for i in range(0, input_sentence_size, batch_length): yield dataset[i: i batch_length][text] # Train tokenizer tokenizer.train_from_iterator( iteratorbatch_iterator(input_sentence_sizeinput_sentence_size), vocab_sizevocab_size, show_progressTrue, ) # Save files to disk tokenizer.save(./norwegian-t5-base/tokenizer.json)T5 詞表大小為 32,000。由于 Unigram 訓(xùn)練需要多輪 EM 迭代文檔提示該步驟最長可達(dá) 120 分鐘是四種模型中耗時(shí)最長的。7.2 創(chuàng)建配置from transformers import T5Config config T5Config.from_pretrained(google/t5-v1_1-base, vocab_sizetokenizer.get_vocab_size()) config.save_pretrained(./norwegian-t5-base)這里vocab_size直接取 tokenizer 實(shí)際訓(xùn)練出的詞表大小保證詞嵌入維度與 tokenizer 嚴(yán)格對齊。7.3 啟動預(yù)訓(xùn)練python run_t5_mlm_flax.py \ --output_dir./norwegian-t5-base \ --model_typet5 \ --config_name./norwegian-t5-base \ --tokenizer_name./norwegian-t5-base \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length512 \ --per_device_train_batch_size32 \ --per_device_eval_batch_size32 \ --adafactor \ --learning_rate0.005 \ --weight_decay0.001 \ --warmup_steps2000 \ --overwrite_output_dir \ --logging_steps500 \ --save_steps10000 \ --eval_steps2500 \ --push_to_hub關(guān)鍵差異--adafactorT5 訓(xùn)練慣例上使用 Adafactor 而非 AdamW這是文檔明確給出的配置也是run_t5_mlm_flax.py中TrainingArguments.adafactor開關(guān)的典型用法max_seq_length512、batch size 32、save_steps10000間隔更大因?yàn)?T5 單步更耗時(shí)。文檔給出的收斂參考3 個(gè) epoch 后 loss 與準(zhǔn)確率收斂到2.36 與 57.0耗時(shí)約 4.5 小時(shí)。訓(xùn)練統(tǒng)計(jì)可直接在 Hub 模型倉庫的 TensorBoard 頁查看。八、BART 去噪語言建模DLM預(yù)訓(xùn)練 BART-base第四部分演示以 BART 論文《BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension》提出的去噪目標(biāo)預(yù)訓(xùn)練facebook/bart-base。BART 同樣是 encoder-decoder 架構(gòu)但與 T5 的 span-masking 不同它的輸入會疊加多種噪聲如 token 刪除、文本旋轉(zhuǎn)、span 掩碼、句序打亂目標(biāo)是讓 decoder 重建完整原文。8.1 訓(xùn)練 tokenizer 與創(chuàng)建配置mkdir ./norwegian-bart-basetokenizer 訓(xùn)練與 RoBERTa 完全一致ByteLevel BPEvocab_size50265代碼可復(fù)用 5.1 節(jié)示例。配置創(chuàng)建from transformers import BartConfig config BartConfig.from_pretrained(facebook/bart-base, vocab_size50265) config.save_pretrained(./norwegian-bart-base)8.2 啟動預(yù)訓(xùn)練python run_bart_dlm_flax.py \ --output_dir./norwegian-bart-base \ --config_name./norwegian-bart-base \ --tokenizer_name./norwegian-bart-base \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length1024 \ --per_device_train_batch_size32 \ --per_device_eval_batch_size32 \ --learning_rate1e-4 \ --warmup_steps2000 \ --overwrite_output_dir \ --logging_steps500 \ --save_steps2000 \ --eval_steps2000 \ --push_to_hub注意 BART 的max_seq_length1024是四種模型中最長的上下文窗口。文檔給出的收斂參考3 個(gè) epoch 后 loss 與準(zhǔn)確率收斂到1.36 與 0.77耗時(shí) 6 小時(shí)以內(nèi)。九、訓(xùn)練參數(shù)全景表來自腳本源碼run_mlm_flax.py、run_clm_flax.py、run_t5_mlm_flax.py中三組 dataclass 定義了全部命令行參數(shù)。下表匯總默認(rèn)值與含義除特別說明外三個(gè)腳本一致TrainingArguments參數(shù)默認(rèn)值說明--output_dir必填預(yù)測結(jié)果與 checkpoint 輸出目錄--overwrite_output_dirFalse輸出目錄非空時(shí)覆蓋指向已有 checkpoint 目錄可續(xù)訓(xùn)--do_train/--do_evalFalse是否執(zhí)行訓(xùn)練 / 評估--per_device_train_batch_size8每個(gè) GPU/TPU 核/CPU 的訓(xùn)練 batch--per_device_eval_batch_size8每個(gè)設(shè)備的評估 batch--learning_rate5e-5初始學(xué)習(xí)率--weight_decay0.0AdamW 權(quán)重衰減系數(shù)--adam_beta10.9Adam 一階矩指數(shù)衰減--adam_beta20.999Adam 二階矩指數(shù)衰減--adam_epsilon1e-8Adam epsilon--adafactorFalse用 Adafactor 替代 AdamW--num_train_epochs3.0訓(xùn)練總 epoch 數(shù)--warmup_steps0線性 warmup 步數(shù)--logging_steps500每 X 步記錄一次日志--save_steps500每 X 步保存一次 checkpoint--eval_stepsNone每 X 步執(zhí)行一次評估--seed42全局隨機(jī)種子--push_to_hubFalse訓(xùn)練后是否上傳模型到 Hub--hub_model_id/--hub_tokenNoneHub 倉庫名 / 認(rèn)證 token--gradient_checkpointingFalse梯度檢查點(diǎn)省顯存換速度僅 MLM 腳本支持ModelArguments參數(shù)默認(rèn)值說明--model_name_or_pathNone初始化權(quán)重的 checkpoint不傳則從零訓(xùn)練--model_typeNone從零訓(xùn)練時(shí)指定模型類型從FLAX_MODEL_FOR_*_LM_MAPPING鍵中取--config_name/--tokenizer_nameNone與模型不同的 config / tokenizer 名稱或路徑--cache_dirNone預(yù)訓(xùn)練模型緩存目錄--use_fast_tokenizerTrue使用 fast tokenizer基于 tokenizers 庫--dtypefloat32權(quán)重初始化與訓(xùn)練精度float32/float16/bfloat16--use_auth_tokenFalse訪問私有模型所需的 Hub tokenDataTrainingArguments以 MLM 腳本為例參數(shù)默認(rèn)值說明--dataset_name/--dataset_config_nameNone通過 datasets 庫加載的數(shù)據(jù)集名與子配置名--train_file/--validation_fileNone本地 csv/json/txt 訓(xùn)練與驗(yàn)證文件--validation_split_percentage5無驗(yàn)證集時(shí)從訓(xùn)練集切出的驗(yàn)證比例--max_seq_lengthNone最大序列長度超過則截?cái)嗄J(rèn)取模型上限--preprocessing_num_workersNone預(yù)處理并行進(jìn)程數(shù)--mlm_probability0.15MLM 掩碼比例--pad_to_max_lengthFalse是否全部 padding 到max_seq_length--line_by_lineFalse是否按行切分序列否則拼接全部文本再切塊實(shí)戰(zhàn)提示腳本同時(shí)支持把參數(shù)寫成 JSON 文件傳入python run_mlm_flax.py args.json源碼通過HfArgumentParser.parse_json_file支持便于在集群上固化實(shí)驗(yàn)配置。十、運(yùn)行時(shí)評估Flax vs PyTorch/XLA vs PyTorch文檔對同一個(gè) MLM 任務(wù)在三種后端/硬件組合下做了耗時(shí)對比用于說明 Flax 在 TPU 上的編譯與并行優(yōu)勢。所有實(shí)驗(yàn)均運(yùn)行在 Google Cloud Platform 上任務(wù)TPU v3-8FlaxTPU v3-8PyTorch/XLA8 GPUPyTorchMLM15h32m23h46m44h14m說明GPU 實(shí)驗(yàn)除 JAX 變換外未做額外優(yōu)化且以全精度 fp32 運(yùn)行TPU v3-8 指 4 塊芯片共 8 個(gè) TPU 核8 GPU 指 8 塊獨(dú)立 GPU 芯片。10.1 用 PyTorch/XLA 在 TPUv3-8 上復(fù)現(xiàn)同一實(shí)驗(yàn)前提是已按 5.1/5.2 節(jié)在norwegian-roberta-base目錄備好 tokenizer 與配置。在示例目錄下建立符號鏈接把 PyTorch 版run_mlm.py與xla_spawn.py引入當(dāng)前工作目錄ln -s ~/transformers/examples/pytorch/language-modeling/run_mlm.py ./ ln -s ~/transformers/examples/pytorch/xla_spawn.py ./設(shè)置環(huán)境變量export XRT_TPU_CONFIGlocalservice;0;localhost:51011 unset LD_PRELOAD export NUM_TPUS8 export TOKENIZERS_PARALLELISM0 export MODEL_DIR./norwegian-roberta-base mkdir -p ${MODEL_DIR}啟動訓(xùn)練python3 xla_spawn.py --num_cores ${NUM_TPUS} run_mlm.py --output_dir./runs \ --model_typeroberta \ --config_name${MODEL_DIR} \ --tokenizer_name${MODEL_DIR} \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length128 \ --weight_decay0.01 \ --per_device_train_batch_size128 \ --per_device_eval_batch_size128 \ --learning_rate3e-4 \ --warmup_steps1000 \ --overwrite_output_dir \ --num_train_epochs18 \ --adam_beta10.9 \ --adam_beta20.98 \ --do_train \ --do_eval \ --logging_steps500 \ --evaluation_strategyepoch \ --report_totensorboard \ --save_strategyno10.2 用 PyTorch 在 8 張 V100 GPU 上復(fù)現(xiàn)同一實(shí)驗(yàn)由于單張 V100 顯存限制最大 batch 只能到 32Flax 版為 128必須引入gradient_accumulation彌補(bǔ)全局 batchln -s ~/transformers/examples/pytorch/language-modeling/run_mlm.py ./export NUM_GPUS8 export TOKENIZERS_PARALLELISM0 export MODEL_DIR./norwegian-roberta-base mkdir -p ${MODEL_DIR}python3 -m torch.distributed.launch --nproc_per_node ${NUM_GPUS} run_mlm.py \ --output_dir${MODEL_DIR} \ --model_typeroberta \ --config_name${MODEL_DIR} \ --tokenizer_name${MODEL_DIR} \ --dataset_nameoscar \ --dataset_config_nameunshuffled_deduplicated_no \ --max_seq_length128 \ --weight_decay0.01 \ --per_device_train_batch_size32 \ --per_device_eval_batch_size32 \ --gradient_accumulation4 \ --learning_rate3e-4 \ --warmup_steps1000 \ --overwrite_output_dir \ --num_train_epochs18 \ --adam_beta10.9 \ --adam_beta20.98 \ --do_train \ --do_eval \ --logging_steps500 \ --evaluation_strategysteps \ --report_totensorboard \ --save_strategynoper_device_train_batch_size32乘以gradient_accumulation4再乘以 8 卡全局有效 batch 與 Flax 版128×8等價(jià)保證對比公平。十一、從預(yù)訓(xùn)練到推理與 FlexGen 的關(guān)系完成上述任一預(yù)訓(xùn)練流程后得到的 checkpoint 即可被加載用于推理/微調(diào)。這與倉庫主項(xiàng)目 FlexGen 構(gòu)成完整的上下游鏈條FlexGen 專注于讓大語言模型在單 GPU 上以吞吐優(yōu)先的方式高效推理見倉庫根目錄 README.md 與flexgen/下的實(shí)現(xiàn)而本文所述 Flax 示例則負(fù)責(zé)模型從零到有的預(yù)訓(xùn)練環(huán)節(jié)。對于想在生產(chǎn)環(huán)境落地大模型的讀者可以先用本文流程訓(xùn)練領(lǐng)域語料模型再借助 FlexGen 的思路做推理側(cè)的吞吐優(yōu)化兩種能力互為補(bǔ)充覆蓋了從訓(xùn)練到服務(wù)的完整生命周期。十二、常見問題與調(diào)參建議收斂不佳時(shí)優(yōu)先檢查什么確認(rèn)vocab_size與 tokenizer 實(shí)際詞表一致T5 場景務(wù)必用tokenizer.get_vocab_size()確認(rèn)max_seq_length不超過模型支持上限源碼中max_seq_length min(data_args.max_seq_length, tokenizer.model_max_length)顯存不足降低per_device_*_batch_size或開啟--gradient_checkpointing僅run_mlm_flax.py提供該開關(guān)PyTorch 對比實(shí)驗(yàn)則用gradient_accumulation補(bǔ)償多機(jī)多卡Flax 腳本依賴jax.device_count()自動感知設(shè)備數(shù)訓(xùn)練 batch 會按設(shè)備數(shù)自動放大無需手工配置分布式啟動器自定義語料不需要 Hub 數(shù)據(jù)集時(shí)直接用--train_file/--validation_file傳入 csv/json/txt 文件即可腳本會自動按text列或第一列讀取復(fù)現(xiàn)性--seed默認(rèn) 42腳本在模型初始化與數(shù)據(jù)洗牌處均注入該種子。本文所有腳本、配置與依賴均可在當(dāng)前倉庫benchmark/third_party/transformers/examples/flax/language-modeling/目錄下找到并直接運(yùn)行建議結(jié)合源碼閱讀本文以獲得對 JAX/Flax 訓(xùn)練管線最完整的理解。贊分享推理引擎大模型【免費(fèi)下載鏈接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.項(xiàng)目地址https://gitcode.com/gh_mirrors/fl/FlexGen點(diǎn)擊查看免費(fèi)下載相關(guān)推薦使用 Ray Train JaxTrainer 從零分布式訓(xùn)練 GPT-2JAX/Flax 規(guī)?；?xùn)練實(shí)戰(zhàn)指南使用 Ray Train JaxTrainer 從零分布式訓(xùn)練 GPT 2JAX/Flax 規(guī)?；?xùn)練實(shí)戰(zhàn)指南 預(yù)計(jì)閱讀與實(shí)操時(shí)間 約 15 分鐘 本指南以人工智能分布式訓(xùn)練強(qiáng)化學(xué)習(xí)任務(wù)調(diào)度模型推理服務(wù)后端使用 SpeechBrain 訓(xùn)練 Switchboard 語言模型Transformer LM 從零訓(xùn)練與預(yù)訓(xùn)練微調(diào)實(shí)戰(zhàn)指南使用 SpeechBrain 訓(xùn)練 Switchboard 語言模型Transformer LM 從零訓(xùn)練與預(yù)訓(xùn)練微調(diào)實(shí)戰(zhàn)指南 本篇技術(shù)指南以 SpeechB人工智能深度學(xué)習(xí)語音音頻NLP預(yù)訓(xùn)練Fairseq 多語言 RoBERTa 預(yù)訓(xùn)練實(shí)戰(zhàn)指南從多語語料預(yù)處理到 masked LM 訓(xùn)練Fairseq 多語言 RoBERTa 預(yù)訓(xùn)練實(shí)戰(zhàn)指南從多語語料預(yù)處理到 masked LM 訓(xùn)練 多語言 RoBERTaMultilingual RoBE人工智能深度學(xué)習(xí)預(yù)訓(xùn)練NLP語音上一篇FastGPT 工作流變量替換 CPU 阻塞優(yōu)化從變量驅(qū)動到模板驅(qū)動的調(diào)度熱路徑重構(gòu)下一篇LeetCode 179 Largest Number 的 Go 題解自定義比較器快排與 ab、ba 拼接比較創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考