使用星雲提升檢查點速度並降低成本

若要啟用與 PyTorch 型訓練腳本的完整 Nebula 相容性，請視需要修改您的定型腳本。

1. 首先，匯入必要的 nebulaml 套件：

     # Import the Nebula package for fast-checkpointing 
     import nebulaml as nm
   

2. 若要初始化 Nebula，請在 中main()呼叫 函nm.init()式，如下所示：

     # Initialize Nebula with variables that helps Nebula to know where and how often to save your checkpoints
     persistent_storage_path="/tmp/test",
     nm.init(persistent_storage_path, persistent_time_interval=2)
   

3. 若要儲存檢查點，請取代原始 torch.save() 語句，以將檢查點儲存為星雲。 請確定您的檢查點實例開頭為 「global_step」，例如 「global_step500」 或 「global_step1000」：

   checkpoint = nm.Checkpoint('global_step500')  
   checkpoint.save('<CKPT_NAME>', model)  
   

   注意

   <'CKPT_TAG_NAME'> 是檢查點的唯一標識碼。 標籤通常是步驟數目、epoch 編號或任何使用者定義的名稱。 選擇性的選擇性 <'NUM_OF_FILES'> 參數會指定您要為這個標籤儲存的狀態號碼。

4. 載入最新的有效檢查點，如下所示：

   latest_ckpt = nm.get_latest_checkpoint()
   p0 = latest_ckpt.load(<'CKPT_NAME'>)
   

   由於檢查點或快照集可能包含許多檔案，因此您可以依名稱載入一或多個檔案。 使用最新的檢查點，定型狀態可以還原到最後一個檢查點所儲存的狀態。

   其他 API 可以處理檢查點管理

   • 列出所有檢查點
   • 取得最新的檢查點

   # Managing checkpoints
   ## List all checkpoints
   ckpts = nm.list_checkpoints()
   ## Get Latest checkpoint path
   latest_ckpt_path = nm.get_latest_checkpoint_path("checkpoint", persisted_storage_path)
   

如果您在組ds_config.json態檔中啟用星雲，則以 DeepSpeed 為基礎的定型腳本 （>=0.7.3） 可以使用星雲，如下所示：

"nebula": {
      "enabled": true,
      "persistent_storage_path": "<YOUR STORAGE PATH>",
      "persistent_time_interval": 100,
      "num_of_version_in_retention": 2,
      "enable_nebula_load": true
  }

此 JSON 代碼段函式的運作方式類似 函式 nebulaml.init() 。

使用 ds_config.json 檔案組態進行初始化會啟用星雲，進而啟用檢查點儲存。 原始的 DeepSpeed 儲存方法，搭配模型檢查點 API model_engine.save_checkpoint()會自動使用星雲。 這個儲存方法可避免需要修改程序代碼。

PyTorch Lightning （Nebula 支援在使用定型器時自動支援版本 >=1.5.0） 檢查點。 由於您通常會使用自定義行為來儲存檢查點，以進行精細的控制，PyTorch Lightning 提供兩種方式來儲存檢查點：使用 ModelCheckpoint()條件式儲存 ，並使用 手動儲存 trainer.save_checkpoint()。 這些技術適用於 PyTorch （>=0.15.0） 定型腳本。

如果您使用 ModelCheckpoint 來有條件地儲存檢查點，您可以使用 NebulaCallback 而不是 ModelCheckpoint 進行初始化。

# import Nebula package
import nebulaml as nm

# define NebulaCallback
config_params = dict()
config_params["persistent_storage_path"] = "<YOUR STORAGE PATH>"
config_params["persistent_time_interval"] = 10

nebula_checkpoint_callback = nm.NebulaCallback(
   ****, # Original ModelCheckpoint params
   config_params=config_params, # customize the config of init nebula
)

接下來，將 新增 nm.NebulaCheckpointIO() 為外掛程式至 ， Trainer並修改 trainer.save_checkpoint() 記憶體參數，如下所示：

trainer = Trainer(plugins=[nm.NebulaCheckpointIO()],   # add NebulaCheckpointIO as a plugin
                  callbacks=[nebula_checkpoint_callback]) # use NebulaCallback as a plugin

如果您使用 trainer.save_checkpoint() 手動儲存檢查點，您可以在 NebulaCheckpointIO 中使用 Trainer外掛程式，並在 中 trainer.save_checkpoint() 修改記憶體參數，如下所示：

# import Nebula package
import nebulaml as nm

# initialize Nebula
nm.init(persistent_storage_path=<YOUR STORAGE PATH>) 

trainer = Trainer(plugins=[nm.NebulaCheckpointIO()])  # add NebulaCheckpointIO as a plugin

# Saving checkpoints
storage_options = {}
storage_options['is_best'] = True
storage_options['persist_path'] = "/tmp/tier3/checkpoint"

trainer.save_checkpoint("example.ckpt",
  storage_options=storage_options, # customize the config of Nebula saving checkpoint
)

載入檢查點

我們會載入與 PyTorch 和 PyTorch Lightning 一致的檢查點。 唯一的修改會指定記憶體路徑。

若要載入最新的檢查點， MyLightningModule.load_from_checkpoint() 仍可運作，如下所示：

  persistent_path = "/tmp/tier3/checkpoint"
  latest_ckpt_path = nebulaml.get_latest_checkpoint_path("checkpoint", persist_path)
  model = MyLightningModule.load_from_checkpoint(latest_ckpt_path)

如果您在 nebula_checkpoint_callback 中使用 Trainer()，則原始未變更的文稿仍可運作：

  trainer = Trainer(
      default_root_dir=tmpdir, 
      max_steps=100,
      plugins=[NebulaCheckpointIO()],
      callbacks=[nebula_checkpoint_callback],
  )
  trainer.fit(model, ckpt_path="/path/example.ckpt")

下一步

• 使用 MLflow 追蹤 ML 實驗和模型
• 記錄和檢視計量