workload_generator/generate_deepspeed_stage1_2_workload.py [124:132]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
                        )
                    )
                    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]),
                            ),
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



workload_generator/generate_deepspeed_stage3_workload.py [148:158]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
                    )
                )

                # weight grad
                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]),
                        ),
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -