本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。

處理中復原和無檢查點訓練

HyperPod 無檢查點訓練使用模型備援來啟用容錯訓練。核心原則是模型和最佳化工具狀態完全複寫到多個節點群組中，並在每個群組中同步複寫權重更新和最佳化工具狀態變更。當失敗時，運作狀態良好的複本會完成其最佳化工具步驟，並將更新後的模型/最佳化工具狀態傳輸至復原複本。

此模型備援型方法可啟用多種錯誤處理機制：

無檢查點組態

以下是無檢查點訓練的核心程式碼片段。

from hyperpod_checkpointless_training.inprocess.train_utils import wait_rank
    wait_rank()
      
def main():
    @HPWrapper(
        health_check=CudaHealthCheck(),
        hp_api_factory=HPAgentK8sAPIFactory(),
        abort_timeout=60.0,
        checkpoint_manager=PEFTCheckpointManager(enable_offload=True),
        abort=CheckpointlessAbortManager.get_default_checkpointless_abort(),
        finalize=CheckpointlessFinalizeCleanup(),
    )
    def run_main(cfg, caller: Optional[HPCallWrapper] = None):
        ...
        trainer = Trainer(
            strategy=CheckpointlessMegatronStrategy(...,
                num_distributed_optimizer_instances=2),
            callbacks=[..., CheckpointlessCallback(...)],
            )
        trainer.fresume = resume
        trainer._checkpoint_connector = CheckpointlessCompatibleConnector(trainer)
        trainer.wrapper = caller

請參閱範例，將無檢查點訓練功能新增至程式碼。

概念

本節介紹無檢查點訓練概念。Amazon SageMaker HyperPod 上的無檢查點訓練支援處理中復原。此 API 界面遵循與 NVRx APIs類似的格式。

概念 - 可重新執行的程式碼區塊 (RCB)

當失敗發生時，運作狀態良好的程序仍然有效，但一部分的程式碼必須重新執行，才能復原訓練狀態和 Python 堆疊。可重新執行的程式碼區塊 (RCB) 是在故障復原期間重新執行的特定程式碼區段。在下列範例中，RDB 包含整個訓練指令碼 （即 main() 下的所有項目），這表示每個失敗復原都會重新啟動訓練指令碼，同時保留記憶體內模型和最佳化工具狀態。

概念 - 故障控制

當無檢查點訓練期間發生失敗時，故障控制器模組會收到通知。此故障控制器包含下列元件：

概念 - 模型備援

大型模型訓練通常需要夠大的資料平行大小，才能有效率地訓練模型。在 PyTorch DDP 和 Horovod 等傳統資料平行處理中，模型會完全複寫。更進階的碎片資料平行處理技術，例如 DeepSpeed ZeRO 最佳化工具與 FSDP 也支援混合碎片模式，允許碎片群組內的模型/最佳化工具狀態碎片，以及跨複寫群組完整複寫。NeMo 也透過允許備援的引數 num_distributed_optimizer_instances 具有此混合碎片功能。

不過，新增備援表示模型不會在整個叢集中完全碎片化，導致更高的裝置記憶體使用量。備援記憶體的數量會根據使用者實作的特定模型碎片技術而有所不同。低精度模型權重、漸層和啟用記憶體不會受到影響，因為它們會透過模型平行處理進行碎片處理。高精度的主模型權重/梯度和最佳化工具狀態會受到影響。新增一個備援模型複本可增加裝置記憶體用量，大約等同於一個 DCP 檢查點大小。

混合碎片會將整個 DP 群組的集合分成相對較小的集合。之前，整個 DP 群組都有一個減少散射和一個全集。在混合碎片之後，減少散射只會在每個模型複本內執行，而且跨模型複本群組會有全部減少。所有集合也會在每個模型複本內執行。因此，整個通訊磁碟區大致保持不變，但集體執行的群組較小，因此我們預期會有更好的延遲。

概念 - 失敗和重新啟動類型

下表記錄不同的故障類型和相關聯的復原機制。無檢查點訓練會先透過程序內復原嘗試失敗復原，然後重新啟動程序層級。只有在發生災難性故障 （例如，多個節點同時故障） 時，才會回到任務層級重新啟動。

概念 - 最佳化工具步驟的原子鎖定保護

模型執行分為三個階段：向前傳播、向後傳播和最佳化工具步驟。復原行為會根據失敗時間而有所不同：

此策略可確保絕對不會捨棄已完成的最佳化工具更新，有助於縮短故障復原時間。

無檢查點訓練流程圖

下列步驟概述故障偵測和無檢查點復原程序：

API 參考

wait_rank

hyperpod_checkpointless_training.inprocess.train_utils.wait_rank()

等待並從 HyperPod 擷取排名資訊，然後使用分散式訓練變數更新目前的程序環境。

此函數會為分散式訓練取得正確的排名指派和環境變數。它可確保每個程序都能取得分散式訓練任務中其角色的適當組態。

參數

無

傳回值

無

行為

環境變數

函數會讀取下列環境變數：

引發

範例

from hyperpod_checkpointless_training.inprocess.train_utils import wait_rank  

# Call before initializing distributed training  
wait_rank()  

# Now environment variables are properly set for this rank  
import torch.distributed as dist  
dist.init_process_group(backend='nccl')

備註

HPWrapper

class hyperpod_checkpointless_training.inprocess.wrap.HPWrapper(  
    *,  
    abort=Compose(HPAbortTorchDistributed()),  
    finalize=None,  
    health_check=None,  
    hp_api_factory=None,  
    abort_timeout=None,  
    enabled=True,  
    trace_file_path=None,  
    async_raise_before_abort=True,  
    early_abort_communicator=False,  
    checkpoint_manager=None,  
    check_memory_status=True)

Python 函數包裝函式，可在 HyperPod 無檢查點訓練中啟用可重新執行程式碼區塊 (RCB) 的重新啟動功能。

此包裝函式透過監控訓練執行，並在發生故障時跨分散式程序協調重新啟動，提供容錯能力和自動復原功能。它使用內容管理員方法，而不是裝飾項目，在整個訓練生命週期中維護全域資源。

參數

方法

def __call__(self, fn)

包裝函數以啟用重新啟動功能。

參數：

傳回：

範例

from hyperpod_checkpointless_training.nemo_plugins.checkpoint_manager import CheckpointManager  
from hyperpod_checkpointless_training.nemo_plugins.patches import patch_megatron_optimizer  
from hyperpod_checkpointless_training.nemo_plugins.checkpoint_connector import CheckpointlessCompatibleConnector  
from hyperpod_checkpointless_training.inprocess.train_utils import HPAgentK8sAPIFactory  
from hyperpod_checkpointless_training.inprocess.abort import CheckpointlessFinalizeCleanup, CheckpointlessAbortManager   
      
@HPWrapper(  
    health_check=CudaHealthCheck(),  
    hp_api_factory=HPAgentK8sAPIFactory(),  
    abort_timeout=60.0,  
    checkpoint_manager=CheckpointManager(enable_offload=False),  
    abort=CheckpointlessAbortManager.get_default_checkpointless_abort(),  
    finalize=CheckpointlessFinalizeCleanup(),  
)def training_function():  
    # Your training code here  
    pass

備註

HPCallWrapper

class hyperpod_checkpointless_training.inprocess.wrap.HPCallWrapper(wrapper)

在執行期間監控和管理重新啟動程式碼區塊 (RCB) 的狀態。

此類別處理 RCB 執行的生命週期，包括故障偵測、與重新啟動的其他排名協調，以及清除操作。它可管理分散式同步，並確保所有訓練程序的一致性復原。

參數

Attributes

方法

def initialize_barrier()

遇到 RCB 的例外狀況後，請等待 HyperPod 障礙物同步處理。

def start_hp_fault_handling_thread()

啟動錯誤處理執行緒以監控和協調失敗。

def handle_fn_exception(call_ex)

處理執行函數或 RCB 的例外狀況。

參數：

def restart(term_ex)

執行重新啟動處理常式，包括定案、垃圾收集和運作狀態檢查。

參數：

def launch(fn, *a, **kw)

以適當的例外狀況處理執行 RCB。

參數：

def run(fn, a, kw)

處理重新啟動和障礙同步的主要執行迴圈。

參數：

def shutdown()

關閉錯誤處理和監控執行緒。

備註

CudaHealthCheck

class hyperpod_checkpointless_training.inprocess.health_check.CudaHealthCheck(timeout=datetime.timedelta(seconds=30))

確保目前程序的 CUDA 內容在無檢查點訓練復原期間處於良好狀態。

此運作狀態檢查會與 GPU 同步，以確認 CUDA 內容在訓練失敗後未損毀。它會執行 GPU 同步操作，以偵測任何可能無法成功恢復訓練的問題。運作狀態檢查會在分散式群組銷毀並完成之後執行。

參數

方法

__call__(state, train_ex=None)

執行 CUDA 運作狀態檢查來驗證 GPU 內容完整性。

參數：

傳回：

引發：

狀態保留：如果所有檢查都通過，則傳回原始狀態和例外狀況不變

範例

import datetime  
from hyperpod_checkpointless_training.inprocess.health_check import CudaHealthCheck  
from hyperpod_checkpointless_training.inprocess.wrap import HPWrapper  
  
# Create CUDA health check with custom timeout  
cuda_health_check = CudaHealthCheck(  
    timeout=datetime.timedelta(seconds=60)  
)  
  
# Use with HPWrapper for fault-tolerant training  
@HPWrapper(  
    health_check=cuda_health_check,  
    enabled=True  
)  
def training_function():  
    # Your training code here  
    pass

備註

HPAgentK8sAPIFactory

class hyperpod_checkpointless_training.inprocess.train_utils.HPAgentK8sAPIFactory()

用於建立 HPAgentK8sAPI 執行個體的工廠類別，可與 HyperPod 基礎設施通訊，以進行分散式訓練協調。

此工廠提供標準化方法來建立和設定 HPAgentK8sAPI 物件，以處理訓練程序與 HyperPod 控制平面之間的通訊。它會封裝基礎通訊端用戶端和 API 執行個體的建立，確保訓練系統不同部分的組態一致。

方法

__call__()

建立並傳回為 HyperPod 通訊設定的 HPAgentK8sAPI 執行個體。

傳回：

範例

from hyperpod_checkpointless_training.inprocess.train_utils import HPAgentK8sAPIFactory  
from hyperpod_checkpointless_training.inprocess.wrap import HPWrapper  
from hyperpod_checkpointless_training.inprocess.health_check import CudaHealthCheck  
  
# Create the factory  
hp_api_factory = HPAgentK8sAPIFactory()  
  
# Use with HPWrapper for fault-tolerant training  
hp_wrapper = HPWrapper(  
    hp_api_factory=hp_api_factory,  
    health_check=CudaHealthCheck(),  
    abort_timeout=60.0,  
    enabled=True  
)  
  
@hp_wrapper  
def training_function():  
    # Your distributed training code here  
    pass

備註

CheckpointManager

class hyperpod_checkpointless_training.nemo_plugins.checkpoint_manager.CheckpointManager(  
    enable_checksum=False,  
    enable_offload=False)

管理記憶體內檢查點和peer-to-peer復原，以在分散式訓練中實現無檢查點容錯能力。

此類別提供 HyperPod 無檢查點訓練的核心功能，包括管理記憶體中的 NeMo 模型檢查點、驗證復原可行性，以及在運作狀態良好和失敗排名之間協調peer-to-peer檢查點轉移。它無需在復原期間使用磁碟 I/O，可大幅縮短復原的平均時間 (MTTR)。

參數

Attributes

方法

save_checkpoint(trainer)

將 NeMo 模型檢查點儲存在記憶體中，以實現潛在的無檢查點復原。

參數：

備註：

delete_checkpoint()

刪除記憶體內檢查點並執行清除操作。

備註：

try_checkpointless_load(trainer)

從對等排名載入狀態以嘗試無檢查點復原。

參數：

傳回：

備註：

checkpointless_recovery_feasible(trainer, include_checksum_verification=True)

判斷目前失敗案例是否可以進行無檢查點復原。

參數：

傳回：

驗證條件：

store_rng_states()

存放所有隨機數字產生器狀態以進行確定性復原。

備註：

load_rng_states()

還原所有 RNG 狀態以繼續進行確定性復原。

備註：

maybe_offload_checkpoint()

如果啟用卸載，將檢查點從 GPU 卸載到 CPU 記憶體。

備註：

範例

from hyperpod_checkpointless_training.inprocess.wrap import HPWrapper  
from hyperpod_checkpointless_training.nemo_plugins.checkpoint_manager import CheckpointManager  
# Use with HPWrapper for complete fault tolerance  
@HPWrapper(  
    checkpoint_manager=CheckpointManager(),  
    enabled=True  
)  
def training_function():  
    # Training code with automatic checkpointless recovery  
    pass

驗證：使用檢查總和驗證檢查點完整性 （如果已啟用）

備註

PEFTCheckpointManager

class hyperpod_checkpointless_training.nemo_plugins.checkpoint_manager.PEFTCheckpointManager(  
    *args,  
    **kwargs)

使用個別的基礎和轉接器處理來管理 PEFT （參數效率微調） 的檢查點，以最佳化無檢查點復原。

這個專門的檢查點管理器擴展 CheckpointManager，透過將基本模型權重與轉接器參數分開來最佳化 PEFT 工作流程。

參數

從 CheckpointManager 繼承所有參數：

其他屬性

方法

maybe_save_base_model(trainer)

儲存基礎模型權重一次，篩選掉轉接器參數。

參數：

備註：

save_checkpoint(trainer)

將 NeMo PEFT 轉接器模型檢查點儲存在記憶體中，以實現潛在的無檢查點復原。

參數：

備註：

try_base_model_checkpointless_load(trainer)

嘗試 PEFT 基礎模型透過從對等排名載入狀態來加權無檢查點復原。

參數：

傳回：

備註：

try_checkpointless_load(trainer)

嘗試 PEFT 轉接器透過從對等排名載入狀態來加權無檢查點復原。

參數：

傳回：

備註：

is_adapter_key(key)

檢查狀態索引鍵是否屬於轉接器參數。

參數：

傳回：

偵測邏輯：

maybe_offload_checkpoint()

將基本模型權重從 GPU 卸載至 CPU 記憶體。

備註：

備註

範例

from hyperpod_checkpointless_training.inprocess.wrap import HPWrapper  
from hyperpod_checkpointless_training.nemo_plugins.checkpoint_manager import PEFTCheckpointManager  
# Use with HPWrapper for complete fault tolerance  
@HPWrapper(  
    checkpoint_manager=PEFTCheckpointManager(),  
    enabled=True  
)  
def training_function():  
    # Training code with automatic checkpointless recovery  
    pass

CheckpointlessAbortManager

class hyperpod_checkpointless_training.inprocess.abort.CheckpointlessAbortManager()

工廠類別，用於建立和管理無檢查點容錯能力的中止元件組成。

此公用程式類別提供靜態方法來建立、自訂和管理 HyperPod 無檢查點訓練中錯誤處理期間使用的中止元件合成。它簡化了中止序列的組態，以便在故障復原期間處理分散式訓練元件、資料載入器和架構特定資源的清除。

參數

無 （所有方法皆為靜態）

靜態方法

get_default_checkpointless_abort()

取得包含所有標準中止元件的預設中止合成執行個體。

傳回：

預設元件：

create_custom_abort(abort_instances)

建立僅具有指定中止執行個體的自訂中止撰寫。

參數：

傳回：

引發：

override_abort(abort_compose, abort_type, new_abort)

使用新元件取代 Compose 執行個體中的特定中止元件。

參數：

傳回：

引發：

範例

from hyperpod_checkpointless_training.inprocess.wrap import HPWrapper  
from hyperpod_checkpointless_training.nemo_plugins.callbacks import CheckpointlessCallback  
from hyperpod_checkpointless_training.inprocess.abort import CheckpointlessFinalizeCleanup, CheckpointlessAbortManager  
  
# The strategy automatically integrates with HPWrapper  
@HPWrapper(  
    abort=CheckpointlessAbortManager.get_default_checkpointless_abort(),  
    health_check=CudaHealthCheck(),  
    finalize=CheckpointlessFinalizeCleanup(),  
    enabled=True  
)  
def training_function():  
    trainer.fit(...)

備註

CheckpointlessFinalizeCleanup

class hyperpod_checkpointless_training.inprocess.abort.CheckpointlessFinalizeCleanup()

在故障偵測後執行全面清理，以準備在無檢查點訓練期間進行程序內復原。

此最終處理常式會透過銷毀訓練元件參考來執行架構特定的清除操作，包括 Megatron/TransformerEngine 中止、DDP 清除、模組重新載入和記憶體清除。它可確保訓練環境正確重設，以成功進行程序內復原，而無需完全終止程序。

參數

無

Attributes

方法

__call__(*a, **kw)

執行完整的清除操作，以進行程序內復原準備。

參數：

清除操作：

register_attributes(trainer)

註冊訓練器執行個體，以便在清除操作期間使用。

參數：

與 CheckpointlessCallback 整合

from hyperpod_checkpointless_training.nemo_plugins.callbacks import CheckpointlessCallback  
from hyperpod_checkpointless_training.inprocess.wrap import HPWrapper  
  
# The strategy automatically integrates with HPWrapper  
@HPWrapper(  
    ...  
    finalize=CheckpointlessFinalizeCleanup(),   
)  
def training_function():  
    trainer.fit(...)

備註

CheckpointlessMegatronStrategy

class hyperpod_checkpointless_training.nemo_plugins.megatron_strategy.CheckpointlessMegatronStrategy(*args, **kwargs)

NeMo Megatron 策略具有整合式無檢查點復原功能，可用於容錯分散式訓練。

請注意，無檢查點訓練num_distributed_optimizer_instances至少需要 2 個，才能進行最佳化工具複寫。此策略也負責基本屬性註冊和程序群組初始化。

參數

從 MegatronStrategy 繼承所有參數：

Attributes

方法

setup(trainer)

初始化策略，並向培訓人員註冊容錯能力元件。

參數：

設定操作：

setup_distributed()

使用具有字首或無根連線的 TCPStore 初始化程序群組。

load_model_state_dict(checkpoint, strict=True)

具有無檢查點復原相容性的載入模型狀態指示。

參數：

get_wrapper()

取得 HPCallWrapper 執行個體以進行容錯能力協調。

傳回：

is_peft()

透過檢查 PEFT 回呼，檢查訓練組態中是否已啟用 PEFT （參數效率微調）

傳回：

teardown()

覆寫 PyTorch Lightning 原生縮減以委派清除以中止處理常式。

範例

from hyperpod_checkpointless_training.inprocess.wrap import HPWrapper  
  
# The strategy automatically integrates with HPWrapper  
@HPWrapper(  
    checkpoint_manager=checkpoint_manager,  
    enabled=True  
)  
def training_function():  
    trainer = pl.Trainer(strategy=CheckpointlessMegatronStrategy())  
    trainer.fit(model, datamodule)

CheckpointlessCallback

class hyperpod_checkpointless_training.nemo_plugins.callbacks.CheckpointlessCallback(  
    enable_inprocess=False,  
    enable_checkpointless=False,  
    enable_checksum=False,  
    clean_tensor_hook=False,  
    clean_lightning_module=False)

閃電回呼，整合 NeMo 訓練與無檢查點訓練的容錯能力系統。

此回呼會管理處理中復原功能的步驟追蹤、檢查點儲存和參數更新協調。它可做為 PyTorch Lightning 訓練迴圈與 HyperPod 無檢查點訓練機制之間的主要整合點，協調整個訓練生命週期的容錯能力操作。

參數

Attributes

方法

get_wrapper_from_trainer(trainer)

從訓練人員取得 HPCallWrapper 執行個體，以進行容錯能力協調。

參數：

傳回：

on_train_batch_start(trainer, pl_module, batch, batch_idx, *args, **kwargs)

在每個訓練批次開始時呼叫 ，以管理步驟追蹤和復原。

參數：

on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)

在每個訓練批次結束時釋放參數更新鎖定。

參數：

備註：

get_peft_callback(trainer)

從訓練人員的回呼清單中擷取 PEFT 回呼。

參數：

傳回：

_try_adapter_checkpointless_restore(trainer, params_to_save)

嘗試無檢查點還原 PEFT 轉接器參數。

參數：

備註：

範例

from hyperpod_checkpointless_training.nemo_plugins.callbacks import CheckpointlessCallback  
from hyperpod_checkpointless_training.nemo_plugins.checkpoint_manager import CheckpointManager  
import pytorch_lightning as pl  
  
# Create checkpoint manager  
checkpoint_manager = CheckpointManager(  
    enable_checksum=True,  
    enable_offload=True  
)  
  
# Create checkpointless callback with full fault tolerance  
checkpointless_callback = CheckpointlessCallback(  
    enable_inprocess=True,  
    enable_checkpointless=True,  
    enable_checksum=True,  
    clean_tensor_hook=True,  
    clean_lightning_module=True  
)  
  
# Use with PyTorch Lightning trainer  
trainer = pl.Trainer(  
    callbacks=[checkpointless_callback],  
    strategy=CheckpointlessMegatronStrategy()  
)  
  
# Training with fault tolerance  
trainer.fit(model, datamodule=data_module)

記憶體管理

CheckpointlessCompatibleConnector

class hyperpod_checkpointless_training.nemo_plugins.checkpoint_connector.CheckpointlessCompatibleConnector()

PyTorch Lightning 檢查點連接器，整合無檢查點復原與傳統的磁碟型檢查點載入。

此連接器擴展 PyTorch Lightning 的 _CheckpointConnector，以提供無檢查點復原與標準檢查點還原之間的無縫整合。它會先嘗試無檢查點復原，然後在無檢查點復原不可行或失敗時回到磁碟型檢查點載入。

參數

從 _CheckpointConnector 繼承所有參數

方法

resume_start(checkpoint_path=None)

嘗試以無檢查點復原優先順序預先載入檢查點。

參數：

resume_end()

完成檢查點載入程序並執行載入後操作。

備註

CheckpointlessAutoResume

class hyperpod_checkpointless_training.nemo_plugins.resume.CheckpointlessAutoResume()

使用延遲設定擴展 NeMo 的 AutoResume，以在檢查點路徑解析之前啟用無檢查點復原驗證。

此類別實作兩階段初始化策略，允許在回到傳統磁碟型檢查點載入之前進行無檢查點復原驗證。它會有條件地延遲 AutoResume 設定，以防止過早檢查點路徑解析，讓 CheckpointManager 先驗證無檢查點peer-to-peer復原是否可行。

參數

從 AutoResume 繼承所有參數

方法

setup(trainer, model=None, force_setup=False)

有條件地延遲 AutoResume 設定，以啟用無檢查點復原驗證。

參數：

範例

from hyperpod_checkpointless_training.nemo_plugins.resume import CheckpointlessAutoResume  
from hyperpod_checkpointless_training.nemo_plugins.megatron_strategy import CheckpointlessMegatronStrategy  
import pytorch_lightning as pl  
  
# Create trainer with checkpointless auto-resume  
trainer = pl.Trainer(  
    strategy=CheckpointlessMegatronStrategy(),  
    resume=CheckpointlessAutoResume()  
)

備註