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

記憶體映射的資料載入器

另一個重新啟動額外負荷來自資料載入：訓練叢集在資料載入器初始化、從遠端檔案系統下載資料，並分批處理時保持閒置狀態。

為了解決這個問題，我們推出了記憶體映射 DataLoader(MMAP) 資料載入器，它在持久性記憶體中快取預先擷取的批次，確保它們即使在錯誤引發的重新啟動之後仍然可用。此方法可省去資料載入器設定時間，並可讓訓練使用快取批次立即恢復，同時資料載入器在背景中重新初始化和擷取後續資料。資料快取位於需要訓練資料的每個排名，並維護兩種類型的批次：最近用於訓練的已取用批次，以及準備好立即使用的預先擷取批次。

MMAP 資料載入器提供下列兩個功能：

使用快取，訓練任務將受益於：

MMAP 組態

若要使用 MMAP，只需將原始資料模組傳入 MMAPDataModule

data_module=MMAPDataModule(
    data_module=MY_DATA_MODULE(...),
    mmap_config=CacheResumeMMAPConfig(
        cache_dir=self.cfg.mmap.cache_dir,
        checkpoint_frequency=self.cfg.mmap.checkpoint_frequency),
)

CacheResumeMMAPConfig：MMAP Dataloader 參數控制快取目錄位置、大小限制和資料擷取委派。根據預設，每個節點只有 TP 排名 0 會從來源擷取資料，而相同資料複寫群組中的其他排名會從共用快取讀取，消除備援傳輸。

MMAPDataModule：它會包裝原始資料模組，並傳回 mmap 資料載入器以進行訓練和驗證。

請參閱啟用 MMAP 的範例。

API 參考

CacheResumeMMAPConfig

class hyperpod_checkpointless_training.dataloader.config.CacheResumeMMAPConfig(
  cache_dir='/dev/shm/pdl_cache',
  prefetch_length=10,
  val_prefetch_length=10,
  lookback_length=2,
  checkpoint_frequency=None,
  model_parallel_group=None,
  enable_batch_encryption=False)

HyperPod 無檢查點訓練中快取-恢復記憶體映射 (MMAP) 資料載入器功能的組態類別。

此組態可透過快取和預先擷取功能實現高效率的資料載入，允許透過在記憶體映射檔案中維護快取的資料批次，在失敗後快速恢復訓練。

參數

方法

create(dataloader_init_callable,
    parallel_state_util,
   step,
    is_data_loading_rank,
   create_model_parallel_group_callable,
    name='Train',
   is_val=False,
   cached_len=0)

建立並傳回設定的 MMAP 資料載入器執行個體。

參數

傳回 CacheResumePrefetchedDataLoader或 CacheResumeReadDataLoader – 設定的 MMAP 資料載入器執行個體

ValueError 如果步驟參數為 ，則引發 None。

範例

from hyperpod_checkpointless_training.dataloader.config import CacheResumeMMAPConfig

# Create configuration
config = CacheResumeMMAPConfig(
    cache_dir="/tmp/training_cache",
    prefetch_length=20,
    checkpoint_frequency=100,
    enable_batch_encryption=False
)

# Create dataloader
dataloader = config.create(
    dataloader_init_callable=my_dataloader_init,
    parallel_state_util=parallel_util,
    step=current_step,
    is_data_loading_rank=lambda: rank == 0,
    create_model_parallel_group_callable=create_mp_group,
    name="TrainingData"
)

備註

MMAPDataModule

class hyperpod_checkpointless_training.dataloader.mmap_data_module.MMAPDataModule(  
    data_module,  
    mmap_config,  
    parallel_state_util=MegatronParallelStateUtil(),  
    is_data_loading_rank=None)

PyTorch Lightning DataModule 包裝函式，可將記憶體映射 (MMAP) 資料載入功能套用至現有的 DataModules，以進行無檢查點訓練。

此類別會包裝現有的 PyTorch Lightning DataModule，並使用 MMAP 功能增強它，以便在訓練失敗期間實現有效的資料快取和快速復原。它維持與原始 DataModule 介面的相容性，同時新增無檢查點訓練功能。

Parameters

Attributes

方法

setup(stage=None)

為指定的訓練階段設定基礎資料模組。

train_dataloader()

使用 MMAP 包裝建立訓練 DataLoader。

傳回：DataLoader – MMAP 包裝的訓練 DataLoader 搭配快取和預先擷取功能

val_dataloader()

使用 MMAP 包裝建立驗證 DataLoader。

傳回：DataLoader – 使用快取功能的 MMAP 包裝驗證 DataLoader

test_dataloader()

如果基礎資料模組支援，請建立測試 DataLoader。

傳回：DataLoader 或無 – 從基礎資料模組測試 DataLoader，如果不支援則測試無

predict_dataloader()

如果基礎資料模組支援，請建立預測 DataLoader。

傳回：DataLoader 或 None – 從基礎資料模組預測 DataLoader，如果不支援則為 None

load_checkpoint(checkpoint)

載入檢查點資訊，以從特定步驟繼續訓練。

get_underlying_data_module()

取得基礎包裝的資料模組。

傳回：pl.LightningDataModule – 包裝的原始資料模組

state_dict()

取得 MMAP DataModule 的狀態字典以進行檢查點。

傳回：dict – 包含快取資料載入器長度的字典

load_state_dict(state_dict)

載入狀態字典以還原 MMAP DataModule 狀態。

屬性

data_sampler

將基礎資料模組的資料取樣器公開至 NeMo 架構。

傳回：物件或無 – 基礎資料模組的資料取樣器

範例

from hyperpod_checkpointless_training.dataloader.mmap_data_module import MMAPDataModule  
from hyperpod_checkpointless_training.dataloader.config import CacheResumeMMAPConfig  
from my_project import MyLLMDataModule  

# Create MMAP configuration  
mmap_config = CacheResumeMMAPConfig(  
    cache_dir="/tmp/training_cache",  
    prefetch_length=20,  
    checkpoint_frequency=100  
)  

# Create original data module  
original_data_module = MyLLMDataModule(  
    data_path="/path/to/data",  
    batch_size=32  
)  

# Wrap with MMAP capabilities  
mmap_data_module = MMAPDataModule(  
    data_module=original_data_module,  
    mmap_config=mmap_config  
)  

# Use in PyTorch Lightning Trainer  
trainer = pl.Trainer()  
trainer.fit(model, data=mmap_data_module)  

# Resume from checkpoint  
checkpoint = {"global_step": 1000}  
mmap_data_module.load_checkpoint(checkpoint)

備註