def backward()

in workload_generator/generate_deepspeed_stage1_2_workload.py [0:0]

• 29 lines of code
• 5 McCabe index (conditional complexity)

    def backward(self):
        self.current_op = "backward"
        for param in self.all_params[::-1]:
            if param.numel() + self.num_in_reduce_bucket > self.max_reduce_bucket_size:
                self._reduce_ipg_grads()
            self.reduce_bucket.append(param)
            self.num_in_reduce_bucket += param.numel()
            if self.compute_enable:
                if param.get_shape()[-1] != 1:
                    self.workload.append(
                        LogItem(
                            comm_type=CommType.computation,
                            msg_size=(
                                (self.batch_size, self.seq_len, param.get_shape()[0]),
                                (param.get_shape()[0], param.get_shape()[1]),
                            ),
                            stage=f"{self.current_op}.computation",
                        )
                    )
                    self.workload.append(
                        LogItem(
                            comm_type=CommType.computation,
                            msg_size=(
                                (param.get_shape()[0], self.batch_size * self.seq_len),
                                (self.batch_size * self.seq_len, param.get_shape()[1]),
                            ),
                            stage=f"{self.current_op}.computation",
                        )
                    )