fix kunlunxin t5_small 1x1 training error (#315)

FlagOpen · Nov 6, 2023 · c506656 · c506656
1 parent 3593f9e
commit c506656
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/training/benchmarks/t5_small/pytorch/train/trainer.py b/training/benchmarks/t5_small/pytorch/train/trainer.py
@@ -44,11 +44,11 @@ def init(self, train_dataloader, eval_dataloader):
         self.optimizer = create_optimizer(self.model, self.config)
         self.lr_scheduler = create_scheduler(self.optimizer, train_dataloader,
                                              self.config)
-
-        self.accelerator = Accelerator()
-        self.model, self.optimizer, train_dataloader, eval_dataloader, self.lr_scheduler = self.accelerator.prepare(
-            self.model, self.optimizer, train_dataloader, eval_dataloader,
-            self.lr_scheduler)
+        if self.config.distributed:
+            self.accelerator = Accelerator()
+            self.model, self.optimizer, train_dataloader, eval_dataloader, self.lr_scheduler = self.accelerator.prepare(
+                self.model, self.optimizer, train_dataloader, eval_dataloader,
+                self.lr_scheduler)
 
         return train_dataloader, eval_dataloader
 
@@ -78,7 +78,10 @@ def train_one_epoch(self, train_dataloader, eval_dataloader):
             outputs = model(**batch)
             loss = outputs.loss
 
-            self.accelerator.backward(loss)
+            if self.config.distributed:
+                self.accelerator.backward(loss)
+            else:
+                loss.backward()
             optimizer.step()
             self.lr_scheduler.step()
             optimizer.zero_grad()