Upload run_cloud_training.py with huggingface_hub

run_cloud_training.py

@@ -27,7 +27,47 @@ from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelFor
 from transformers.data.data_collator import DataCollatorMixin
 from peft import LoraConfig
 from unsloth import FastLanguageModel
-import deepspeed
+
+# Try to import deepspeed and install mpi4py if needed
+try:
+    import deepspeed
+except ImportError:
+    logger.error("DeepSpeed not found. Installing deepspeed...")
+    try:
+        import subprocess
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "deepspeed"])
+        import deepspeed
+        logger.info("DeepSpeed installed successfully")
+    except Exception as e:
+        logger.error(f"Failed to install DeepSpeed: {e}")
+        logger.error("Will continue without DeepSpeed")
+
+# Check for mpi4py which is required by DeepSpeed
+try:
+    import mpi4py
+    logger.info(f"mpi4py is available (version: {mpi4py.__version__})")
+except ImportError:
+    logger.warning("mpi4py not found. Installing mpi4py which is required for DeepSpeed...")
+    try:
+        import subprocess
+        # First try to install OpenMPI if on a Linux system
+        if sys.platform.startswith('linux'):
+            try:
+                logger.info("Attempting to install OpenMPI system dependencies...")
+                subprocess.check_call(["apt-get", "update", "-y"])
+                subprocess.check_call(["apt-get", "install", "-y", "libopenmpi-dev", "openmpi-bin"])
+                logger.info("OpenMPI installed successfully")
+            except Exception as e:
+                logger.warning(f"Failed to install OpenMPI system dependencies: {e}")
+                logger.warning("Will try to install mpi4py anyway")
+        
+        # Now install mpi4py
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "mpi4py>=3.1.4"])
+        import mpi4py
+        logger.info(f"mpi4py installed successfully (version: {mpi4py.__version__})")
+    except Exception as e:
+        logger.error(f"Failed to install mpi4py: {e}")
+        logger.error("DeepSpeed may not work correctly without mpi4py")
 
 # Disable all attention optimizations that might cause issues
 os.environ["TRANSFORMERS_NO_FLASH_ATTENTION"] = "1"
@@ -581,31 +621,51 @@ def train(config_path, dataset_name, output_dir):
         if deepspeed_config:
             logger.info("DeepSpeed configuration found - enabling DeepSpeed for distributed training")
             
-            # Create a temporary DeepSpeed config file
-            ds_config_path = os.path.join(output_dir, "ds_config_temp.json")
-            
-            # Update DeepSpeed config with dynamic values
-            if isinstance(deepspeed_config.get("train_micro_batch_size_per_gpu"), str) and deepspeed_config.get("train_micro_batch_size_per_gpu") == "auto":
-                deepspeed_config["train_micro_batch_size_per_gpu"] = per_device_train_batch_size
+            # Check if mpi4py is available
+            mpi4py_available = False
+            try:
+                import mpi4py
+                mpi4py_available = True
+            except ImportError:
+                logger.error("mpi4py is required for DeepSpeed but not available")
+                logger.error("Will continue without DeepSpeed")
                 
-            if isinstance(deepspeed_config.get("train_batch_size"), str) and deepspeed_config.get("train_batch_size") == "auto":
-                deepspeed_config["train_batch_size"] = per_device_train_batch_size * gpu_count
-            
-            # Write the DeepSpeed config to a file
-            with open(ds_config_path, 'w') as f:
-                json.dump(deepspeed_config, f, indent=2)
-            
-            logger.info(f"Created DeepSpeed config at {ds_config_path}")
-            logger.info(f"DeepSpeed ZeRO Stage: {deepspeed_config.get('zero_optimization', {}).get('stage', 'Not specified')}")
-            
-            # Enable CPU offloading if configured
-            if deepspeed_config.get("zero_optimization", {}).get("offload_optimizer", {}).get("device") == "cpu":
-                logger.info("DeepSpeed CPU offloading enabled for optimizer states")
-            
-            # Set using_deepspeed flag
-            using_deepspeed = True
+            if mpi4py_available:
+                try:
+                    # Create a temporary DeepSpeed config file
+                    ds_config_path = os.path.join(output_dir, "ds_config_temp.json")
+                    
+                    # Update DeepSpeed config with dynamic values
+                    if isinstance(deepspeed_config.get("train_micro_batch_size_per_gpu"), str) and deepspeed_config.get("train_micro_batch_size_per_gpu") == "auto":
+                        deepspeed_config["train_micro_batch_size_per_gpu"] = per_device_train_batch_size
+                        
+                    if isinstance(deepspeed_config.get("train_batch_size"), str) and deepspeed_config.get("train_batch_size") == "auto":
+                        deepspeed_config["train_batch_size"] = per_device_train_batch_size * gpu_count
+                    
+                    # Write the DeepSpeed config to a file
+                    with open(ds_config_path, 'w') as f:
+                        json.dump(deepspeed_config, f, indent=2)
+                    
+                    logger.info(f"Created DeepSpeed config at {ds_config_path}")
+                    logger.info(f"DeepSpeed ZeRO Stage: {deepspeed_config.get('zero_optimization', {}).get('stage', 'Not specified')}")
+                    
+                    # Enable CPU offloading if configured
+                    if deepspeed_config.get("zero_optimization", {}).get("offload_optimizer", {}).get("device") == "cpu":
+                        logger.info("DeepSpeed CPU offloading enabled for optimizer states")
+                    
+                    # Set using_deepspeed flag
+                    using_deepspeed = True
+                except Exception as e:
+                    logger.error(f"Failed to initialize DeepSpeed: {e}")
+                    logger.error("Will continue without DeepSpeed")
+                    ds_config_path = None
+                    using_deepspeed = False
+            else:
+                ds_config_path = None
+                using_deepspeed = False
         else:
             logger.warning("No DeepSpeed configuration found - continuing without DeepSpeed")
+            ds_config_path = None
             using_deepspeed = False
         
         # Initialize model with our safe loading function
@@ -649,6 +709,7 @@ def train(config_path, dataset_name, output_dir):
             reports = ["none"]
             logger.warning("No reporting backends available - training metrics won't be logged")
         
+        # Prepare training arguments
         training_args_dict = {
             "output_dir": output_dir,
             "num_train_epochs": training_config.get("num_train_epochs", 3),
@@ -671,9 +732,20 @@ def train(config_path, dataset_name, output_dir):
             "remove_unused_columns": False,
             "seed": 42,
             "dataloader_num_workers": 4,  # Use multiple workers for data loading
-            "deepspeed": ds_config_path  # Add DeepSpeed config path if available
         }
-    
+        
+        # Add DeepSpeed config if available and mpi4py is installed
+        if using_deepspeed and ds_config_path is not None:
+            logger.info("Using DeepSpeed for training")
+            training_args_dict["deepspeed"] = ds_config_path
+        else:
+            logger.info("Not using DeepSpeed - falling back to standard distributed training")
+            # If DeepSpeed is not available, ensure we're still using distributed training efficiently
+            if gpu_count > 1:
+                logger.info(f"Using standard distributed training with {gpu_count} GPUs")
+                training_args_dict["local_rank"] = int(os.environ.get("LOCAL_RANK", -1))
+                training_args_dict["gradient_checkpointing"] = True
+        
         # Create TrainingArguments with validated parameters
         training_args = TrainingArguments(**training_args_dict)