训练设置

参数解析

DeepSpeed 使用 argparse 库来为 DeepSpeed 运行时提供命令行配置。使用 deepspeed.add_config_arguments() 将 DeepSpeed 的内置参数添加到您的应用程序的解析器中。

parser = argparse.ArgumentParser(description='My training script.')
parser.add_argument('--local_rank', type=int, default=-1,
                    help='local rank passed from distributed launcher')
# Include DeepSpeed configuration arguments
parser = deepspeed.add_config_arguments(parser)
cmd_args = parser.parse_args()
deepspeed.add_config_arguments(parser)[source]
更新参数解析器以启用 DeepSpeed 命令行参数的解析。

DeepSpeed 参数集包括以下内容：1) –deepspeed：启用 DeepSpeed 的布尔标志 2) –deepspeed_config <json 文件路径>：配置 DeepSpeed 运行时的 json 配置文件的路径。

参数

parser – 参数解析器

返回值

更新的解析器

返回类型

parser

训练初始化

所有使用 DeepSpeed 的训练的入口点是 deepspeed.initialize()。如果分布式后端尚未初始化，则会初始化它。

示例用法

model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
                                                     model=net,
                                                     model_parameters=net.parameters())
deepspeed.initialize(args=None, model: Optional[Module] = None, optimizer: Optional[Union[Optimizer, Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer]] = None, model_parameters: Optional[Module] = None, training_data: Optional[Dataset] = None, lr_scheduler: Optional[Union[_LRScheduler, Callable[[Optimizer], _LRScheduler]]] = None, distributed_port: int = 29500, mpu=None, dist_init_required: Optional[bool] = None, collate_fn=None, config=None, mesh_param=None, config_params=None)[source]

初始化 DeepSpeed 引擎。

参数

  • args – 包含 local_rank 和 deepspeed_config 字段的对象。如果传递了 config，则可选。

  • model – 必需：应用任何包装器之前的 nn.module 类

  • optimizer – 可选：用户定义的优化器或返回优化器对象的 Callable。这将覆盖 DeepSpeed json 配置中的任何优化器定义。

  • model_parameters – 可选：torch.Tensors 或字典的可迭代对象。指定应优化的张量。

  • training_data – 可选：类型为 torch.utils.data.Dataset 的数据集

  • lr_scheduler – 可选：学习率调度器对象或接受优化器并返回调度器对象的 Callable。调度器对象应定义 get_lr()、step()、state_dict() 和 load_state_dict() 方法

  • distributed_port – 可选：主节点 (rank 0) 在分布式训练期间需要使用的空闲端口

  • mpu – 可选：实现 get_{model,data}_parallel_{rank,group,world_size}() 的模型并行单元对象

  • dist_init_required – 可选：None 将在需要时自动初始化 torch 分布式，否则用户可以通过布尔值强制初始化或不初始化。

  • collate_fn – 可选：将样本列表合并以形成张量的小批量。在从地图式数据集进行批量加载时使用。

  • config – 可选：代替需要 args.deepspeed_config，您可以将 deepspeed 配置作为参数传递，作为路径或字典。

  • config_params – 可选：与 config 相同，保留用于向后兼容。

返回值

一个包含 engine、optimizer、training_dataloader、lr_scheduler 的元组

  • engine：DeepSpeed 运行时引擎，它包装客户端模型以进行分布式训练。

  • optimizer：如果用户定义了 optimizer 或如果在 json 配置中指定了优化器，则为包装的优化器，否则为 None。

  • training_dataloader：如果提供了 training_data，则为 DeepSpeed 数据加载器，否则为 None。

  • lr_scheduler：如果传递了用户 lr_scheduler 或如果在 JSON 配置中指定了 lr_scheduler，则为包装的 lr 调度器。否则为 None。

分布式初始化

可选的分布式后端初始化，与 deepspeed.initialize() 分开。在用户想要在调用 deepspeed.initialize() 之前使用 torch 分布式调用时很有用，例如在使用模型并行、流水线并行或某些数据加载器场景时。

deepspeed.init_distributed(dist_backend=None, auto_mpi_discovery=True, distributed_port=29500, verbose=True, timeout=datetime.timedelta(seconds=1800), init_method=None, dist_init_required=None, config=None, rank=-1, world_size=-1)[source]

初始化 dist 后端，如果需要，可能执行 MPI 发现

参数

  • dist_backend – 可选 (str)。torch 分布式后端，例如，nccl、mpi、gloo、hccl

  • 可选 (auto_mpi_discovery) –

  • distributed_port – 可选 (int)。torch 分布式后端端口

  • verbose – 可选 (bool)。详细日志记录

  • timeout – 可选 (timedelta)。对针对进程组执行的操作的超时。默认值为 30 分钟，可以通过环境变量 DEEPSPEED_TIMEOUT 覆盖。

  • init_method – 可选 (string)。Torch 分布式，指定如何初始化进程组的 URL。如果未指定 init_method 或 store，则默认值为“env://”。

  • config – 可选 (dict)。用于设置通信选项 (例如，通信分析) 的 DeepSpeed 配置

  • rank – 可选（int）。当前手动指定的排名。一些 init_method 如“tcp://” 也需要 rank 和 world_size（参见：https://pytorch.ac.cn/docs/stable/distributed.html#tcp-initialization）。

  • world_size – 可选（int）。TCP 或共享文件系统初始化所需的 world_size。