Fix enable_model_cpu_offload problems (#320)

* update

* update

* update

Author: a-r-r-o-w

finetrainers/args.py

@@ -853,8 +853,9 @@ def _validate_dataset_args(args: BaseArgs):
 
 
 def _validate_validation_args(args: BaseArgs):
-    if args.dp_shards > 1 and args.enable_model_cpu_offload:
-        raise ValueError("Model CPU offload is not supported with FSDP at the moment.")
+    if args.enable_model_cpu_offload:
+        if any(x > 1 for x in [args.pp_degree, args.dp_degree, args.dp_shards, args.cp_degree, args.tp_degree]):
+            raise ValueError("Model CPU offload is not supported on multi-GPU at the moment.")
 
 
 def _display_helper_messages(args: argparse.Namespace):

finetrainers/trainer/sft_trainer/trainer.py

@@ -730,6 +730,7 @@ def _validate(self, step: int, final_validation: bool = False) -> None:
 
         parallel_backend.wait_for_everyone()
         if not final_validation:
+            self._move_components_to_device()
             self.transformer.train()
 
     def _evaluate(self) -> None: