fix dit training (#752)

Author: nemonameless

ppdiffusers/examples/class_conditional_image_generation/DiT/diffusion/__init__.py

@@ -16,7 +16,13 @@
 from .dit import DiT
 from .dit_llama import DiT_Llama
 from .respace import SpacedDiffusion, space_timesteps
-from .trainer import LatentDiffusionTrainer, LatentDiffusionAutoTrainer
+from .trainer import LatentDiffusionTrainer
+try:
+    from paddlenlp.trainer.auto_trainer import AutoTrainer
+    from .trainer_auto import LatentDiffusionAutoTrainer
+except:
+    print(f'please install paddlepaddle-gpu>=3.0.0b2 if using auto trainer')
+
 from .trainer_args import (
     DataArguments,
     ModelArguments,

ppdiffusers/examples/class_conditional_image_generation/DiT/diffusion/trainer.py

@@ -23,7 +23,6 @@
 from paddle.distributed import fleet
 from paddle.io import get_worker_info
 from paddlenlp.trainer import Trainer
-from paddlenlp.trainer.auto_trainer import AutoTrainer
 from paddlenlp.trainer.integrations import (
     INTEGRATION_TO_CALLBACK,
     TrainerCallback,
@@ -295,20 +294,6 @@ def __impl__():
 
     return __impl__
 
-class LatentDiffusionAutoTrainer(AutoTrainer):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def _get_meshes_for_loader(self):
-        def _get_mesh(pp_idx=0):
-            return fleet.auto.get_mesh().get_mesh_with_dim("pp")[pp_idx]
-
-        return _get_mesh(0) # label_id is not label
-
-    def _wrap_for_dist_loader(self, train_dataloader):
-        dist_loader = super()._wrap_for_dist_loader(train_dataloader)
-        dist_loader._input_keys = ["latents", "label_id"]
-        return dist_loader
 
 class LatentDiffusionTrainer(Trainer):
     def __init__(self, **kwargs):