DeepSpeed代码阅读笔记之：elasticity

本笔记是DeepSpeed代码阅读的第一篇笔记，本周的主要任务是阅读DeepSpeed python代码中的

DeepSpeed 部署

DeepSpeed 部署的过程如下：

1. 安装cuda与pytorch
2. 按照requirements文件夹安装依赖：

pip install -r requirements/requirements-dev.txt 
pip install -r requirements/requirements.txt
pip install -r requirements/requirements-sparse_attn.txt 
pip install mpi4py
pip install --ignore-installed PyYAML

注意：attn文件里的triton 可能没有1.0版本
mpi4py可能需要通过conda安装
3. 安装DeepSpeed:运行install.sh
注意：deepspeed需要全目录有rw权限

至此，安装已完成，可以使用DeepSpeedExample目录下的例程来测试。

elasticity目录代码

init中说明了本目录下有如下几个文件：

--elasticity
    |-- __init__.py
    |-- elasticity.py
    |-- utils.py
    |-- contants.py
    |-- elastic_agent.py
    |-- config.py

_init_.py

这个代码主要是把目录下的文件中包含外部接口的文件进行引用，需要判断pytorch版本是否大于1.11，只有高于此版本才能使用DSElasticAgent类

elasticity

elasticity.py中给出了几个接口函数用来供外部调用

compute_elastic_config

调用弹性计算的核心代码，在DeepSpeedConfig类中会检查config是否有配置弹性计算，如果有会调用这个函数最终得到总batch——size和根据当前可用GPU数量得到的micro_batch(用于数据并行)（可选）

elasticity 0.1版本和0.2版本分别调用不同的函数来得到final_batch_size。

最后，通过检查micro_batch_size能否在数据并行中和batch_size 匹配(batch_size分在每个GPU上的大小能为micro_batch_size整倍)

感觉这部分代码的去耦合做的很不好

elasticity_enabled

检查ds_config的ELSASTICITY是否启用

ensure_immutable_elastic_config

确保在资源管理器启用的情况下，根据环境变量中给Deepspeed的弹性配置和deepspeed自身的配置来检查是否匹配。

除此之外，elasticity中内部还有关于micro_batch_size和total_batch_size计算实现的代码

_get_compatible_gpus_v01

这个函数主要是得到batch_size和可用gpu数量。首先按照指定的micro_batches得到候选的batch_size。这个过程是给定的mrbs来找出最大合适的batch_size，存储在列表里。
在此之后，通过batch_sized候选列表中，按照偏好（大or小batch）得到最好的合适的（满足GPU数量和偏好）的batch_size。

注意：这里满足GPU数量是指mrbs可以被bs整除，同时在给定最大or最小GPU数量中。

_get_compatible_gpus_v02

在v01函数的基础上，这个函数根据bs的大小和GPU，根据节点的GPU数量来调整mrbs数量。以及根据数据并行dpsz来得到bs=bs*dpsz。

utils

仅负责检查torch版本是否匹配

constant

存储了必需的常数，类似ENABLE，DS最低版本，环境变量名。默认bs等。

elastic_agent

实现了pytorch LocalElasticAgent的子类。

_set_master_addr_port

这个方法检查主节点的地址（端口），如果没有会自动生成

_start_workers

这个方法使用torch distributed的WorkerGroup类作为参数，给每个worker设定必要的环境变量后，给关于本地worker数量的环境变量进行更新。同时指定必要的参数后，启动workers的进程。

_invoke_run

这个方法在worker启动后，每隔一段时间监控当前workerGroup的状态。
可能会遇到节点工作失败的状况，则会选择进行重启worker或者在全部失效时进行报错推出
在遇到节点加入或退出时，会进行记录并重启workers。

config.py

这个文件主要定义了和elasticity相关的错误抛出，以及对config从ds_config到elasticity_config参数的变换和类型检测。

问题

1. 为什么在elasticity文件中有预定义的HCN_LIST
2. DS相比pytorch的Elasticity多了什么功能？