在使用DeepSpeed进行训练时,配置和初始化步骤是至关重要的。以下是关于DeepSpeed训练设置的详细说明:
参数解析
DeepSpeed使用argparse库为DeepSpeed运行时提供命令行配置。开发者可以通过调用deepspeed.add_config_arguments()函数将自己的应用参数解析器中添加DeepSpeed内置的参数。这样就可以在启动脚本时传递有关DeepSpeed的配置信息。例如:
            
            
              python
              
              
            
          
          parser = argparse.ArgumentParser(description='My training script.')
parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher')
# 添加DeepSpeed配置参数
parser = deepspeed.add_config_arguments(parser)
cmd_args = parser.parse_args()DeepSpeed添加的参数包括:
- --deepspeed:布尔标志,用于启用DeepSpeed
- --deepspeed_config <json file path>:指定一个JSON配置文件的路径,用于配置DeepSpeed运行时。
训练初始化
所有使用DeepSpeed的训练流程均始于deepspeed.initialize()函数。该函数会自动初始化分布式后端,如果尚未初始化的话。
            
            
              python
              
              
            
          
          model_engine, optimizer, _, _ = deepspeed.initialize(
    args=cmd_args,
    model=net,
    model_parameters=net.parameters()
)deepspeed.initialize()函数接收一系列参数来初始化DeepSpeed引擎:
- args: 包含- local_rank和- deepspeed_config字段的对象,如果传入了- config参数,则- args不是必需的。
- model: 必需参数,未经任何封装前的原始nn.module类实例。
- optimizer: 可选参数,用户定义的优化器对象或返回优化器对象的可调用函数,这将覆盖DeepSpeed JSON配置文件中定义的优化器。
- model_parameters: 可选参数,一个包含待优化参数的迭代器,通常为模型参数列表或字典。
- training_data: 可选参数,类型为- torch.utils.data.Dataset的训练数据集。
- lr_scheduler: 可选参数,学习率调度器对象或返回调度器对象的可调用函数,该调度器需要有- get_lr(),- step(),- state_dict(), 和- load_state_dict()等方法。
- distributed_port: 可选参数,主机(rank 0)用于分布式训练通信的可用端口。
- mpu: 可选参数,一个实现模型并行相关接口的模型并行单元对象。
- dist_init_required: 可选参数,如果设为None,当需要时会自动初始化torch分布式;否则,用户可通过布尔值控制是否初始化。
- collate_fn: 可选参数,用于将样本列表合并成mini-batch的函数,适用于从map-style数据集中批量加载数据。
- config: 可选参数,可以直接传入DeepSpeed配置文件的路径或字典,替代从命令行通过- --deepspeed_config传入的方式。
- config_params: 可选参数,与- config相同,保留以兼容旧版本。
函数返回一个包含引擎、优化器、训练数据加载器和学习率调度器的元组。
分布式初始化
为了满足某些特定场景的需求(比如模型并行、流水线并行或者特殊的数据加载器场景),可以在调用deepspeed.initialize()之前单独初始化分布式后端,这时可以使用deepspeed.init_distributed()函数:
            
            
              python
              
              
            
          
          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
)此函数允许用户根据需求选择torch分布式后端,设置通信端口、日志详细程度、超时时间以及其他与初始化相关的参数,以便在不同的并行化方案中正确初始化分布式环境。