完整指南:保存與恢復訓練數據管道狀態(tài))
深度學習數據工程【免費下載鏈接】DALIA GPU-accelerated library containing highly optimized building blocks and an execution engine for data processing to accelerate deep learning training and inference applications.項目地址https://gitcode.com/gh_mirrors/da/DALI點擊查看免費下載導讀本文以 NVIDIA DALI 的官方文檔 advanced_topics_checkpointing.rst 為骨架系統(tǒng)講解 DALI 檢查點Checkpointing功能如何保存管道Pipeline當前狀態(tài)到文件并在之后從檢查點恢復使新管道與舊管道產生完全一致的輸出。該能力對可能被中斷的長時間訓練任務尤其有價值。讀完本文你將掌握enable_checkpointing的開啟方式、Pipeline.checkpoint()的保存與checkpoint參數的恢復流程、fn.external_source的部分支持限制以及 TensorFlow 插件中與tf.train.checkpoint的集成方式并結合倉庫源碼了解底層實現原理。什么是 DALI CheckpointingDALI 的檢查點功能允許你把管道Pipeline的當前狀態(tài)保存到一個文件中之后從該檢查點恢復管道恢復后的新管道將產生與舊管道完全相同的輸出。這對于運行時間較長、很可能被中途打斷的訓練任務尤其有用。從實現上看DALI 管道檢查點包含兩類關鍵信息見 checkpoint.h 與文檔說明管道中所有隨機數生成器RNG的狀態(tài)保證恢復后隨機算子如fn.random.uniform生成的隨機序列與中斷前完全一致每個讀取器Reader的進度保證數據讀取從中斷時的 epoch 與迭代位置繼續(xù)而不是從頭開始。在 C 側這一設計體現為Checkpoint類——它是整個管道級狀態(tài)的聚合通過AddOperator(instance_name)為每個算子注冊獨立的OpCheckpoint并用name2id_映射把算子實例名與檢查點條目關聯(lián)起來Checkpoint還持有iteration_id_當前迭代序號以及來自 Python 側的ExternalContextCheckpoint包含pipeline_data與iterator_data見 checkpoint.h。關鍵設計要點檢查點保存的是有狀態(tài)算子的狀態(tài)。那些不維護用戶可觀察狀態(tài)的算子如解碼器、resize、歸一化等在概念上是無狀態(tài)的不會進入檢查點——這一點在 Dynamic 模式檢查點文檔 中有明確說明。完整實操示例可參考官方 notebookPipeline checkpointing notebook動態(tài)模式Dynamic API的等價流程見 Dynamic mode checkpointing。Checkpointing API 使用詳解開啟檢查點enable_checkpointingTrue要啟用檢查點功能在創(chuàng)建管道時將enable_checkpointing設為True。開啟后DALI 會跟蹤每個算子的狀態(tài)以便按需保存。官方文檔明確指出開啟檢查點不應影響性能。pipeline_def(..., enable_checkpointingTrue) def pipeline(): ... p pipeline()在 Python API 中enable_checkpointing是Pipeline.__init__的可選參數默認值為False見 pipeline.py。該參數會一路傳遞到 C 側Python 層在build()時把該參數寫入PipelineParams見 pipeline.pyC 層在管道反序列化/構建時調用this-EnableCheckpointing()并在執(zhí)行器構建階段通過executor_-EnableCheckpointing(checkpointing_enabled())傳遞給執(zhí)行器見 pipeline.cc 與 pipeline.cc在 ProtoBuf 的PipelineDef消息中也保留了一個enable_checkpointing字段默認false見 dali.proto。從執(zhí)行器內部看啟用檢查點后Executor2會為每次迭代的IterationData預先創(chuàng)建Checkpoint對象InitIterationData中if (config_.checkpointing) iter_data-checkpoint CreateCheckpoint(...)見 exec2.cc每個算子節(jié)點執(zhí)行完成后會把自身的狀態(tài)寫入該迭代對應的檢查點見 exec_node_task.cc。這就是按需保存能夠隨時取到最新狀態(tài)的原因。注意shuffle_after_epochTrue的讀取器在啟用檢查點后樣本打亂的方式可能與未啟用時略有不同。原因是啟用檢查點后讀取器必須在每個 epoch 保存一份初始順序的備份以便恢復詳見下文讀取器如何支持檢查點一節(jié)見 file_label_loader.h。保存檢查點Pipeline.checkpoint()保存檢查點需要調用Pipeline.checkpoint()方法它返回一個序列化后的檢查點字符串內部為序列化后的 Protobuf 消息。也可以傳入文件名作為參數DALI 會直接把檢查點寫入該文件文件內容將被覆蓋。for _ in range(iters): output p.run() # 把檢查點寫入文件 checkpoint p.checkpoint() open(checkpoint_file.cpt, wb) # 或者更簡單 checkpoint p.checkpoint(checkpoint_file.cpt)從源碼看checkpoint()的實現流程是見 pipeline.py先調用self.build()若尚未構建確保管道已就緒調用_get_checkpoint()通過b.ExternalContextCheckpoint()把 Python 側的迭代上下文{iter: self._consumer_iter, epoch_idx: self._epoch_idx}JSON 序列化后放入pipeline_data一起打包見 pipeline.py通過self._pipe.GetSerializedCheckpoint(external_ctx_cpt)觸發(fā) C 側序列化若傳入了filename則以二進制寫模式wb把序列化結果寫入文件并返回該字符串。在 C 側Checkpoint::SerializeToProtobuf會遍歷所有算子的OpCheckpoint對每個算子調用op-SerializeCheckpoint(cpt)取回其序列化狀態(tài)再連同external_ctx_cpt一起打包成 Protobuf 消息返回見 checkpoint.cc。序列化格式定義在 dali.proto 的Checkpoint消息中每個OpCheckpoint包含operator_name算子實例名與operator_state算子狀態(tài)字節(jié)流。注意調用Pipeline.checkpoint()可能會引入可觀測的開銷。官方建議不要過于頻繁地調用它。從源碼也可以看出原因每次調用都需要遍歷整張算子圖、逐個算子收集并序列化狀態(tài)且 GPU 算子在保存狀態(tài)時還可能涉及流同步OpCheckpoint::SetOrder(AccessOrder)用于保證異步保存的 GPU 狀態(tài)在主機側可見見 op_checkpoint.h。從檢查點恢復checkpoint參數之后可以從已保存的檢查點恢復管道狀態(tài)。做法是在構造Pipeline時傳入checkpoint參數?;謴秃蟮墓艿缿敭a生與原始管道完全一致的輸出。checkpoint open(checkpoint_file.cpt, rb).read() p_restored pipeline(checkpointcheckpoint)在 Python API 中checkpoint是Pipeline.__init__的另一個可選參數默認值為None見 pipeline.py。其恢復流程為build()過程中調用_restore_state_from_checkpoint()見 pipeline.py若self._checkpoint is not None則調用 C 側self._pipe.RestoreFromSerializedCheckpoint(self._checkpoint)并把is_restored_from_checkpoint置為True見 pipeline.py。該屬性可通過Pipeline.is_restored_from_checkpoint查詢若傳入的檢查點不是合法的 JSON/Protobuf 數據會拋出錯誤提示請確保檢查點是由相同版本的 DALI 創(chuàng)建的。在 C 側Executor2::RestoreFromCheckpoint會遍歷算子圖中的每個算子節(jié)點逐個調用n.op-RestoreState(cpt.GetOpCheckpoint(n.instance_name))恢復狀態(tài)并把執(zhí)行器的迭代計數設為檢查點中保存的iteration_id如果檢查點中的算子狀態(tài)數量多于當前圖中的算子還會拋出檢查點包含多余算子狀態(tài)的運行時錯誤見 exec2.cc。反序列化側的對應實現是Checkpoint::DeserializeFromProtobuf它會按算子名逐一匹配并把狀態(tài)交還給對應算子見 checkpoint.cc。警告恢復時必須保證恢復的管道與原始管道相同即包含相同的算子、相同的參數。用不同管道創(chuàng)建的檢查點去恢復將導致未定義行為undefined behavior。源碼中也能看到對應的防御邏輯RestoreFromCheckpoint要求檢查點中的每個算子名都能在當前圖中找到找不到或多余都會報錯見 exec2.cc。算子層面的檢查點接口從算子基類看DALI 為每個算子定義了 4 個與檢查點相關的虛函數見 operator.hSaveState(OpCheckpoint cpt, AccessOrder order)把算子狀態(tài)保存到檢查點對象RestoreState(const OpCheckpoint cpt)從檢查點恢復算子狀態(tài)SerializeCheckpoint(const OpCheckpoint cpt)把算子狀態(tài)序列化為字符串DeserializeCheckpoint(OpCheckpoint cpt, const std::string data)反序列化并填充檢查點對象。默認實現會調用CheckpointingUnsupportedError()——即該算子未實現檢查點。只有真正有狀態(tài)、且實現了這些接口的算子讀取器、RNG 相關算子等才參與檢查點。單算子級別的測試覆蓋可參考 checkpoint_test.cc其中分別驗證了 CPU-only、GPU-only、混合Mixed三種圖結構以及序列化/反序列化往返CheckpointTest.CPUOnly、GPUOnly、Mixed、Serialize見 checkpoint_test.cc。讀取器如何支持檢查點原理剖析讀取器是檢查點最主要的受益者。DALI 的讀取器由 Loader負責實際讀樣本驅動其檢查點支持在 loader.h 中定義LoaderStateSnapshot保存了讀取器在 epoch 開始時的基礎狀態(tài)——rng隨機數引擎、current_epoch當前 epoch 數和age樣本年齡計數見 loader.h。Loader 構造時通過options.GetArgumentbool(checkpointing)讀取檢查點開關見 loader.h并在Init()階段就保存一份初始快照。具體到FileReader基于file_label_loader.h啟用檢查點且設置了shuffle_after_epoch時PrepareMetadata會先保存一份初始文件順序的備份backup_file_label_entries_見 file_label_loader.h在每個 epoch 的Reset中啟用檢查點時從備份順序重新洗牌file_label_entries_ backup_file_label_entries_洗牌種子由shuffle_after_epoch_seed_ (current_epoch_ 32)推導——每個 epoch 用不同種子因此恢復后仍能保證隨機分布不受影響同時順序可復現見 file_label_loader.hRestoreStateImpl只需恢復current_epoch其余索引狀態(tài)由 loader 基類的快照機制統(tǒng)一處理見 file_label_loader.h。這正好解釋了文檔中的那條注意事項shuffle_after_epochTrue時啟用檢查點后打亂方式可能略有不同——因為可恢復性優(yōu)先打亂順序的推導方式被調整了。另外file_reader_op.cc中shuffle_after_epoch的文檔還提到使用shuffle_after_epoch時不能同時使用stick_to_shard和random_shuffle多 GPU 場景下所有管道實例應使用相同的shuffle_after_epoch_seed以保證全局一致的洗牌見 file_reader_op.cc。External Source 的檢查點支持部分支持fn.external_source算子僅部分支持檢查點。支持的場景只有當source是單參數可調用對象callable且該參數為以下三者之一時檢查點才受支持批次索引batch indexBatchInfoSampleInfo。對于這類source恢復檢查點后查詢會從檢查點中保存的位置繼續(xù)epoch 與迭代都會對齊。從源碼看Python 側_check_checkpointing_support的實現邏輯正是如此只有kind _SourceKind.CALLABLE and has_inputs即可調用且?guī)挡潘阒С謾z查點否則會發(fā)出警告見 pipeline.py。external_source的文檔也明確注明恢復檢查點后單參數可調用 source 的查詢會從檢查點保存的 epoch 和迭代繼續(xù)見 external_source.py。其底層依靠callback_args中基于current_iter/epoch_idx的索引推導見 external_source.py。不支持的場景其他類型的source如無參可調用、迭代器等不支持檢查點。它們的狀態(tài)不會被保存進檢查點恢復后這些 source 會從頭開始。如果與從中間恢復的讀取器搭配使用可能導致數據錯位。官方建議如果你需要使用檢查點推薦把 source 改寫成受支持的單參數可調用形式。例如def my_source(sample_info: SampleInfo): # 依據 sample_info.idx_in_epoch / iteration 返回對應樣本 return data[sample_info.idx_in_epoch] pipe pipeline_def(..., enable_checkpointingTrue)(...)TensorFlow 插件中的檢查點nvidia.dali.plugin.tf.DALIDataset與 TensorFlow 的tf.train.checkpoint機制深度集成——這意味著你可以用 TensorFlow 標準的檢查點 API如tf.train.Checkpoint手動保存/恢復來同時保存 DALI 管道的狀態(tài)與模型權重無需額外的 DALI 專屬代碼路徑。在插件實現中DALIDataset的保存save與恢復restore鉤子分別通過 C API 完成保存時調用daliPipelineGetCheckpoint拿到檢查點句柄再經daliPipelineSerializeCheckpoint序列化以名為checkpoint的 Tensor 寫入 TensorFlow 的檢查點文件見 dali_dataset_op.cc恢復時從檢查點文件讀取checkpointTensor經daliPipelineDeserializeCheckpointdaliPipelineRestoreCheckpoint恢復到管道見 dali_dataset_op.cc。重要限制插件層面DALIDatasetWithInputs暫不支持檢查點。其checkCheckpointingSupport()會直接返回Unimplemented錯誤Checkpointing is not supported for DALI dataset with inputs見 dali_dataset_op.ccGPU 數據集暫不支持檢查點。同樣由checkCheckpointingSupport()拋出 Checkpointing is not supported for DALI GPU dataset見 dali_dataset_op.cc。警告使用 TensorFlow 插件時請確保滿足上述兩個前提非DALIDatasetWithInputs、非 GPU 數據集否則會在保存/恢復階段直接報錯。常見問題與最佳實踐1. 檢查點保存多久調用一次合適checkpoint()有可觀測的開銷需要遍歷整張算子圖、序列化所有算子狀態(tài)不要每輪迭代都調用。合理做法是每隔固定的若干迭代如每 N 個 epoch或結合訓練框架的定期保存如 TensorFlow 的tf.train.CheckpointManager、PyTorch 的torch.utils.checkpoint風格定期存盤調用一次。2. 恢復后的管道輸出一定一致嗎只要恢復的管道與原始管道結構、參數完全相同且 source 是受支持的單參數 callable恢復后的輸出應當與原始管道完全一致——這正是檢查點保存 RNG 狀態(tài)與 reader 進度的意義。反之管道不同、source 不支持則可能產生未定義行為或數據錯位。3. 檢查點文件是什么格式checkpoint()返回的是一個序列化后的 Protobuf 字符串其中包含每個算子的狀態(tài)operator_nameoperator_state以及 Python 側的pipeline_data/iterator_data。請用二進制模式讀寫示例中open(checkpoint_file.cpt, wb)/rb。4. 啟用檢查點影響性能嗎官方文檔明確說明不應有任何影響。從實現看啟用后執(zhí)行器只是為每次迭代額外維護一個Checkpoint對象并在算子執(zhí)行后寫入狀態(tài)見 exec2.cc開銷主要體現在調用checkpoint()主動保存的那一刻。5. 多管道 / 多 GPU 場景如果使用多個管道如base_iterator.py中的多管道迭代器所有管道必須設置相同的enable_checkpointing值否則會拋ValueError見 base_iterator.py。若部分管道從檢查點恢復而部分沒有迭代器會發(fā)出警告并可能出現意外結果見 base_iterator.py??偨YDALI 的檢查點功能為長時間訓練任務提供了可靠的中斷恢復手段開啟pipeline_def(..., enable_checkpointingTrue)保存p.checkpoint()或p.checkpoint(file.cpt)返回序列化 Protobuf 字符串恢復構造時傳入checkpoint...恢復后的管道輸出與原管道完全一致核心內容所有 RNG 狀態(tài) 每個讀取器的進度限制fn.external_source僅支持單參數 callableTensorFlow 插件的DALIDatasetWithInputs與 GPU 數據集暫不支持。結合源碼可以確認底層由執(zhí)行器Executor2逐迭代維護檢查點、每個算子通過SaveState/RestoreState/SerializeCheckpoint/DeserializeCheckpoint四個接口參與狀態(tài)保存與恢復最終以 Protobuf 消息整體序列化。相關實現與測試文件包括pipeline.pyPython APIcheckpoint/_get_checkpoint/_restore_state_from_checkpointcheckpoint.h 與 checkpoint.cc管道級檢查點聚合與序列化exec2.cc執(zhí)行器側保存/恢復loader.h 與 file_label_loader.h讀取器狀態(tài)快照checkpoint_test.ccCPU/GPU/混合圖及序列化往返測試dali_dataset_op.ccTensorFlow 插件集成贊分享深度學習數據工程【免費下載鏈接】DALIA GPU-accelerated library containing highly optimized building blocks and an execution engine for data processing to accelerate deep learning training and inference applications.項目地址https://gitcode.com/gh_mirrors/da/DALI點擊查看免費下載相關推薦verl檢查點管理訓練狀態(tài)保存與恢復verl檢查點管理訓練狀態(tài)保存與恢復 概述 在大規(guī)模語言模型LLM的強化學習訓練過程中verlVolcano Engine Reinforcement人工智能大模型強化學習RLHF分布式訓練微調LeRobot機器人學習3步構建你的第一個AI機器人控制模型LeRobot機器人學習3步構建你的第一個AI機器人控制模型 想不想讓機器人像人一樣學習新技能 你是否曾夢想過讓機械臂學會抓取物體、讓機器人自主完成復雜人工智能機器學習深度學習機器人具身智能強化學習55項功能全面升級HsMod插件讓你的爐石傳說體驗飛升8倍速55項功能全面升級HsMod插件讓你的爐石傳說體驗飛升8倍速 HsMod是一款基于BepInEx框架開發(fā)的爐石傳說游戲增強插件為玩家提供了從游戲性能優(yōu)化到社游戲開發(fā)上一篇在VS Code中使用Windows Subsystem for Linux(WSL)進行開發(fā)下一篇FlyEnv v4.9.7 版本更新優(yōu)化 FTP 服務與 PHP 安全配置創(chuàng)作聲明:本文部分內容由AI輔助生成(AIGC),僅供參考