DeepSpeed ZeRO: Source Code and Execution Flow

1
self.training_dataloader = self.deepspeed_io(training_data)

1
if has_optimizer:
2
    self._configure_optimizer(optimizer, model_parameters)

1
class ZeroStageEnum(int, Enum):
2
    """ Enum class for possible zero stages """
3
    disabled = 0
4
    optimizer_states = 1    # 只缓存 optimizer states，对应 ZeRO-1
5
    gradients = 2           # 缓存 optimizer states + gradients，对应 ZeRO-2
6
    weights = 3             # 缓存 optimizer states + gradients + weights，对应 ZeRO-3
7
    max_stage = 3

1
# ZeRO stage 1 (False) or 2 (True)
2
self.partition_gradients = partition_grads
3
self.zero_stage_string = "ZeRO-2" if partition_grads else "ZeRO-1"

1
self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))]
2
self.partition_count = [dp_size for i in range(len(self.optimizer.param_groups))]

1
# 遍历所有的参数组
2
for i, param_group in enumerate(self.optimizer.param_groups):
3
    # 当前进程在特定数据并行组中的 rank，用于确定当前进程负责哪一部分参数
4
    partition_id = dist.get_rank(group=self.real_dp_process_group[i])
5
    ...
6
    # 这里把需要 train 的参数都放进了 bit16_groups 里
7
    if self.round_robin_gradients:
8
        # 把第 i 个参数组的参数按照轮询的方式，拆到同一个进程组的不同进程里
9
        round_robin_tensors, round_robin_indices = self._round_robin_reorder(self.bit16_groups[i],dist.get_world_size(group=self.real_dp_process_group[i]))
10
    else:
11
        round_robin_tensors = self.bit16_groups[i]
12
        round_robin_indices = list(range(len(self.bit16_groups[i])))

1
# 把 tensor 拍平，然后根据 NCCL 的对齐要求，对齐到边界，多余的补上 torch.zeros(..)
2
flattened_buffer = self.flatten_dense_tensors_aligned(
3
    self.round_robin_bit16_groups[i],
4
    self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i]),
5
    use_cpu_data=True)
6
...
7
# 尽可能等分作切割，每个 partition 里存的是一部分 flat tensor
8
data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i)
9
self.parallel_partitioned_bit16_groups.append(data_parallel_partitions)
10
...
11
# partition_id之前的所有 tensor 的元素数量，求一个 left bound，算一下当前的长度
12
left_boundary = sum([t.numel() for t in data_parallel_partitions[:partition_id]])
13
curr_partition_size = data_parallel_partitions[partition_id].numel()
14

15
# 计算 padding 的大小; 这里是为了确保每个 partition 的大小是 nccl_start_alignment_factor 的倍数
16
if orig_group_numel <= left_boundary:
17
    padding = curr_partition_size
18
elif orig_group_numel < left_boundary + curr_partition_size:
19
    padding = left_boundary + curr_partition_size - orig_group_numel
20
else:
21
    padding = 0
22
self.groups_padding.append(padding)

1
partition_size = len(self.bit16_groups_flat[i]) / dist.get_world_size(group=self.real_dp_process_group[i])
2
# 取到当前的 partition id 分片的 params
3
params_in_partition, params_not_in_partition, first_offset = self.get_partition_info(
4
    self.round_robin_bit16_groups[i], partition_size, partition_id)
5

6
self.partition_size.append(partition_size)
7
self.params_in_partition.append(params_in_partition)
8
self.params_not_in_partition.append(params_not_in_partition)
9
self.first_offset.append(first_offset)

1
# 这里是如何判断overflow的？
2
# 1. 先判断当前的 partition 是否 overflow
3
# 2. 如果当前的 partition 没有 overflow，再判断是否有其他的 partition 溢出
4
# 4. 做一个all reduce, 把所有的 partition 的 overflow 标志位都同步起来
5
def has_overflow(self, partition_gradients=True):
6
    overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial()
7
    overflow_gpu = get_accelerator().ByteTensor([overflow]) if self.cpu_offload else overflow.byte().to(
8
        get_accelerator().current_device_name())
9

10
    if partition_gradients:
11
        '''This will capture overflow across all data parallel and expert parallel process
12
        Since expert parallel process are a subset of data parallel process'''
13
        dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group)
14

15
    # Since each model parallel GPU carries only part of the model,
16
    # make sure overflow flag is synced across all the model parallel GPUs
17
    self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX)
18

19
    overflow = overflow_gpu[0].item()
20
    return bool(overflow)

1
# 释放掉所有不在当前 partition 中的 params 的梯度
2
self.free_grad_in_param_list(self.params_not_in_partition[i])
3

4
# 这里是取到当前 partition 中的 params 的梯度
5
single_grad_partition = self.flatten(self.averaged_gradients[i]).to(self.single_partition_of_fp32_groups[i].dtype)

1
# ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well
2
if self.zero_optimization_partition_gradients():
3
    self.optimizer.overlapping_partition_gradients_reduce_epilogue()
4

5
# Communicate only at gradient accumulation boundaries
6
elif self.is_gradient_accumulation_boundary():
7
    if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states and hasattr(
8
            self.optimizer, 'reduce_gradients'):
9
        self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism)
10
    else:
11
        grads = None
12
        self.buffered_allreduce_fallback(grads=grads, elements_per_buffer=bucket_size)
13
elif self.zenflow:
14
    self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism)

1
# Step 3:- run the optimizer if no offloading
2
self.timers(OPTIMIZER_STEP_TIMER).start()
3
self._optimizer_step(i)
4

5
def _optimizer_step(self, group_no):
6
    original_param_groups = self.optimizer.param_groups
7
    self.optimizer.param_groups = [original_param_groups[group_no]]
8
    ...
9
    self.optimizer.step()
10
    self.optimizer.param_groups = original_param_groups

1
# 这里是 all gather 所有的 params, 然后更新到当前的 partition 中
2
all_gather_dp_groups(groups_flat=self.bit16_groups_flat,
3
                     partitioned_param_groups=self.parallel_partitioned_bit16_groups,
4
                     dp_process_group=self.real_dp_process_group,
5
                     start_alignment_factor=self.nccl_start_alignment_factor,
6
                     allgather_bucket_size=self.allgather_bucket_size)

1
def backward(self, loss, retain_graph=False, scale_wrt_gas=True):
2
    ...
3
    # 这里是直接diaoyong调用 optimizer 的 backward 方法
4
    self._do_optimizer_backward(loss, retain_graph)
5
    # 这里是对 backward() 后的数据做处理
6
    self._backward_epilogue()
7

8
def _backward_epilogue(self):
9
    self.allreduce_gradients()
10

11
def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE):
12
    ...
13
    # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well
14
    if self.zero_optimization_partition_gradients():
15
        self.optimizer.overlapping_partition_gradients_reduce_epilogue()
16

17
def independent_gradient_partition_epilogue(self):
18
    ...
19
    self.reduce_ipg_grads()
20
    ...

1
buckets = {}
2
for i, (dst, bucket_offset, numel) in enumerate(rank_and_offsets):
3
    grad_slice = tensor.narrow(0, int(bucket_offset), int(numel))
4
    bucket_key = real_dp_process_group[i] if self.use_multi_rank_bucket_allreduce else (
5
        dst, real_dp_process_group[i])
6
    if bucket_key not in buckets:
7
        buckets[bucket_key] = []
8
    if self.use_multi_rank_bucket_allreduce:
9
        buckets[bucket_key].append((dst, grad_slice))
10
    else:
11
        buckets[bucket_key].append(grad_slice)

1
for bucket_key in buckets:
2
    self.allreduce_no_retain(buckets[bucket_key],
3
                             communication_data_type,
4
                             numel_per_bucket=self.reduce_bucket_size,
5
                             rank=dst,
6
                             divide=False,
7
                             process_group=process_group)

1
# Pre forward hook
2
self.forward_hooks.append(module.register_forward_pre_hook(_pre_forward_module_hook))
3
# Post forward hook
4
self.forward_hooks.append(module.register_forward_hook(_post_forward_module_hook))

1
class PartitionedParameterCoordinator:
2
    FORWARD_FETCH_SUBMIT = 'forward_fetch_submit'
3
    FORWARD_FETCH_WAIT = 'forward_fetch_wait'
4
    FORWARD_PREFETCH_SUBMIT = 'forward_prefetch_submit'
5
    BACKWARD_FETCH_SUBMIT = 'backward_fetch_submit'
6
    BACKWARD_FETCH_WAIT = 'backward_fetch_wait'
7
    BACKWARD_PREFETCH_SUBMIT = 'backward_prefetch_submit'
8
    FORWARD_ALL_GATHER = 'forward_all_gather'
9
    BACKWARD_ALL_GATHER = 'backward_all_gather'
10
    """Handles partitioning and gathering of parameters."""

1
class ZeroParamStatus(Enum):
2
    # parameters are fully present and ready for use on all processes
3
    AVAILABLE = 1
4
    # parameters are either partitioned or remote in some or all process
5
    NOT_AVAILABLE = 2
6
    # parameters are being gathered.
7
    INFLIGHT = 3

1
def fetch_sub_module(self, current_submodule: Module, forward: bool) -> None:
2
    ...
3
    params_to_fetch = set(iter_params(current_submodule, recurse=z3_leaf_module(current_submodule)))
4
    fetch_numel = sum([p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE])
5

6
    if fetch_numel > 0:
7
        ...
8
        self.__all_gather_params(params_to_fetch, forward)

1
def release_sub_module(self, submodule: Module, forward=False) -> None:
2
    """release the parameters of a sub module, assuming they meet conditions to
3
    be released."""
4
    #print_rank_0(f"release_sub_module {'fwd' if forward else 'bwd'}: {debug_module2name_id(submodule)}", force=False)
5
    params_to_release = (self.__params_to_release(submodule, self.__step_id) if self.is_complete_trace() else set(
6
        p.ds_id for p in iter_params(submodule, recurse=z3_leaf_module(submodule))))
7

8
    free_data = not z3_leaf_module(submodule) or not self.fast_sharding_for_leaf_module
9
    if not free_data:
10
        # wait for the computation to finish and launch as early as possible.
11
        empty_buffer = torch.empty(1, device=get_accelerator().current_device())
12

13
    for param in iter_params(submodule, recurse=z3_leaf_module(submodule)):
14
        param.ds_active_sub_modules.discard(submodule.ds_id)
15
        if param.ds_id in params_to_release and not param.is_external_param:
16
            self.__release_param(param, free_data)
17
        if not free_data:
18
            if param.ds_id in params_to_release and not param.is_external_param:
19
                # empty buffer ensures that all computations are complete
20
                param.data = empty_buffer