计算世界大小示例

这是一个最小的“hello world”风格示例应用程序，它使用 PyTorch Distributed 来计算世界大小。它不进行 ML 训练，但它确实初始化了进程组并执行了单个集体操作（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)