(self, task, i_task, epochs)
| 61 | |
| 62 | |
| 63 | def train_one_task(self, task, i_task, epochs): |
| 64 | if self.args.local_rank == -1: |
| 65 | device = torch.device("cuda") |
| 66 | else: |
| 67 | torch.cuda.set_device(self.args.local_rank) |
| 68 | device = torch.device("cuda", self.args.local_rank) |
| 69 | |
| 70 | #### TRAIN #### |
| 71 | train_dataloader = self.train_task_list[task] |
| 72 | eval_dataloader = self.eval_task_list[task] |
| 73 | total_steps = epochs * len(train_dataloader) |
| 74 | progress_bar = tqdm(total=total_steps, leave=True, disable=(self.args.global_rank != 0)) |
| 75 | for epoch in range(epochs): |
| 76 | print_rank_0( |
| 77 | f"Beginning of Epoch {epoch+1}/{epochs}, Total Micro Batches {len(train_dataloader)}", |
| 78 | self.args.global_rank) |
| 79 | self.model.train() |
| 80 | |
| 81 | for step, batch in enumerate(train_dataloader): |
| 82 | del batch['sources'] |
| 83 | batch = to_device(batch, device) |
| 84 | outputs = self.model(**batch, use_cache=False) |
| 85 | loss = outputs.loss |
| 86 | # Update the description to include current step and loss, if needed |
| 87 | if self.args.global_rank == 0: |
| 88 | # Update the progress bar |
| 89 | progress_bar.update(1) |
| 90 | description = f"Epoch {epoch+1}, Step {step}, Loss: {loss.item():.4f}" |
| 91 | progress_bar.set_description(description, refresh=False) |
| 92 | |
| 93 | self.model.backward(loss) |
| 94 | # Correct gradient accumulation steps are handled withing the deepspeed engine's backward call. |
| 95 | self.model.step() |
| 96 | |
| 97 | |
| 98 | # Evaluate perplexity on the validation set. |
| 99 | # print_rank_0( |
| 100 | # f"***** Evaluating perplexity, Epoch {epoch+1}/{epochs} *****", |
| 101 | # self.args.global_rank) |
| 102 | # perplexity = self.perplexity_evaluation(eval_dataloader, device) |
| 103 | # print_rank_0(f"ppl: {perplexity}", self.args.global_rank) |
| 104 | # self.model.tput_timer.update_epoch_count() |
| 105 | |
| 106 | |
| 107 | def train_continual(self): |
no test coverage detected