| | |
| | | # Multiply world_size because DistributedDataParallel |
| | | # automatically normalizes the gradient by world_size. |
| | | loss *= self.world_size |
| | | # loss *= self.world_size |
| | | # Scale the loss since we're not updating for every mini-batch |
| | | loss = loss / accum_grad |
| | | |
| | |
| | | optim.zero_grad(set_to_none=True) |
| | | total_time = f"{(time.perf_counter() - time5)/accum_grad:0.3f}" |
| | | time5 = time.perf_counter() |
| | | |
| | | speed_stats["optim_time"] = f"{time5 - time4:0.3f}" |
| | | |
| | | speed_stats["total_time"] = total_time |
| | |
| | | f"data_slice: {data_split_i}/{data_split_num}, " |
| | | f"step_in_slice: {batch_idx + 1}/{batch_num_epoch}, step_in_epoch: {step_in_epoch}, total step: {self.batch_total}, " |
| | | f"(loss_avg_rank: {loss:.3f}), " |
| | | f"(loss_avg_epoch: {loss_avg_epoch:.3f}), " |
| | | f"(ppl_avg_epoch: {math.exp(loss_avg_epoch):.3e}), " |
| | | f"(acc_avg_epoch: {acc_avg_epoch:.3f}), " |
| | | f"(loss_avg_slice: {loss_avg_epoch:.3f}), " |
| | | f"(ppl_avg_slice: {math.exp(loss_avg_epoch):.3e}), " |
| | | f"(acc_avg_slice: {acc_avg_epoch:.3f}), " |
| | | f"(lr: {lr:.3e}), " |
| | | f"{[(k, round(v.detach().cpu().item(), 3)) for k, v in stats.items()]}, " |
| | | f"{speed_stats}, " |