【AutoParallel】Add strategy with more options (#8114)

* add strategy * polish
PaddlePaddle · Mar 14, 2024 · c406d90 · c406d90
1 parent 880d2ea
commit c406d90
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/paddlenlp/trainer/auto_trainer.py b/paddlenlp/trainer/auto_trainer.py
@@ -112,8 +112,10 @@ def _wrap_for_auto(self, model, train_dataloader):
         dist_loader = self._wrap_for_dist_loader(train_dataloader)
 
         if self.args.to_static:
+            unified_strategy = dist.Strategy()
+            unified_strategy._from_legacy_strategy(self.args.strategy)
             return (
-                dist.to_static(model, dist_loader, self.criterion, self.optimizer, strategy=self.args.strategy),
+                dist.to_static(model, dist_loader, self.criterion, self.optimizer, strategy=unified_strategy),
                 dist_loader,
             )
         else: