計算世界大小範例

這是一個最小的「hello world」風格範例應用程式，它使用 PyTorch Distributed 來計算世界大小。它沒有進行機器學習訓練，但它初始化了處理程序群組並執行了一個單一集體操作（all_reduce），這足以驗證基礎架構和排程器設定。

儘管這個應用程式很簡單，但實際的 compute_world_size() 函式被拆分到一個單獨的子模組（.module.util.compute_world_size）中，作為工作區修補邏輯的端到端測試，該邏輯通常會對整個專案目錄而不是單一檔案進行差異修補。此應用程式還使用 Hydra 設定檔，作為如何在使用 TorchX 啟動的應用程式中使用 Hydra 設定檔的說明性範例。

使用 dist.ddp 內建元件執行它，以作為驗證應用程式，以確保堆疊已正確設定，以進行更嚴重的分散式訓練作業。

import hydra
from omegaconf import DictConfig, OmegaConf
from torch.distributed.elastic.multiprocessing.errors import record
from torchx.examples.apps.compute_world_size.module.util import compute_world_size


@record
def run(cfg: DictConfig) -> None:
    print(OmegaConf.to_yaml(cfg))

    if cfg.main.throws:
        raise RuntimeError(f"raising error because cfg.main.throws={cfg.main.throws}")
    compute_world_size(cfg)


if __name__ == "__main__":
    # use compose API to make this compatible with ipython notebooks
    # need to initialize the config directory as a module to make it
    # not depends on rel path (PWD) or abs path (torchx install dir)
    # see: https://hydra.cc/docs/advanced/jupyter_notebooks/
    with hydra.initialize_config_module(
        config_module="torchx.examples.apps.compute_world_size.config"
    ):
        cfg: DictConfig = hydra.compose(config_name="defaults")
        run(cfg)