Upload run_cloud_training.py with huggingface_hub

run_cloud_training.py

@@ -28,46 +28,32 @@ from transformers.data.data_collator import DataCollatorMixin
 from peft import LoraConfig
 from unsloth import FastLanguageModel
 
-# 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")
+# Set DeepSpeed environment variables to disable MPI
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "9994"
+os.environ["RANK"] = "0"
+os.environ["LOCAL_RANK"] = "0"
+os.environ["WORLD_SIZE"] = "1"
 
-# Check for mpi4py which is required by DeepSpeed
+# Try to import deepspeed, install mpi4py if needed
 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 deepspeed
+except ImportError as e:
+    if "mpi4py" in str(e):
+        logger.warning("mpi4py not found, installing...")
         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")
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "mpi4py"])
+            import deepspeed
+            logger.info("Successfully installed mpi4py and imported deepspeed")
+        except Exception as install_error:
+            logger.warning(f"Failed to install mpi4py: {install_error}")
+            logger.warning("Continuing without DeepSpeed MPI support")
+            # Set a flag to disable DeepSpeed later
+            os.environ["DISABLE_DEEPSPEED_MPI"] = "1"
+    else:
+        logger.error(f"Failed to import deepspeed: {e}")
+        raise
 
 # Disable all attention optimizations that might cause issues
 os.environ["TRANSFORMERS_NO_FLASH_ATTENTION"] = "1"
@@ -616,53 +602,42 @@ def train(config_path, dataset_name, output_dir):
         per_device_train_batch_size = 4 if gpu_count >= 4 else 2
         logger.info(f"Using batch size: {per_device_train_batch_size} per device (effective batch size: {per_device_train_batch_size * gpu_count * training_config.get('gradient_accumulation_steps', 4)})")
         
-        # Check if DeepSpeed config is available
+        # Check if DeepSpeed config is available and if MPI is disabled
         deepspeed_config = config.get("deepspeed_config", None)
-        if deepspeed_config:
+        if deepspeed_config and os.environ.get("DISABLE_DEEPSPEED_MPI", "0") != "1":
             logger.info("DeepSpeed configuration found - enabling DeepSpeed for distributed training")
             
-            # 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")
+            # 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 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
+            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
+            
+            # Ensure communication backend is set to avoid MPI
+            if "communication_data_type" not in deepspeed_config:
+                deepspeed_config["communication_data_type"] = "fp16"
+            
+            # 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
+        elif os.environ.get("DISABLE_DEEPSPEED_MPI", "0") == "1":
+            logger.warning("DeepSpeed MPI support is disabled due to missing mpi4py. Continuing without DeepSpeed.")
+            ds_config_path = None
+            using_deepspeed = False
         else:
             logger.warning("No DeepSpeed configuration found - continuing without DeepSpeed")
             ds_config_path = None
@@ -709,7 +684,6 @@ 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),
@@ -734,20 +708,23 @@ def train(config_path, dataset_name, output_dir):
             "dataloader_num_workers": 4,  # Use multiple workers for data loading
         }
         
-        # 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")
+        # Add DeepSpeed config path if available and enabled
+        if using_deepspeed and ds_config_path:
+            logger.info("Adding DeepSpeed configuration to training arguments")
             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
+            logger.info("DeepSpeed is disabled - using standard distributed training")
         
         # Create TrainingArguments with validated parameters
-        training_args = TrainingArguments(**training_args_dict)
+        try:
+            training_args = TrainingArguments(**training_args_dict)
+        except Exception as e:
+            logger.error(f"Failed to create training arguments with DeepSpeed: {e}")
+            if "deepspeed" in training_args_dict:
+                logger.warning("Removing DeepSpeed configuration and trying again")
+                del training_args_dict["deepspeed"]
+                training_args = TrainingArguments(**training_args_dict)
+                using_deepspeed = False
         
         # Create trainer with pre-tokenized collator
         trainer = Trainer(