Upload run_cloud_training.py with huggingface_hub

run_cloud_training.py

@@ -568,8 +568,40 @@ def train(config_path, dataset_name, output_dir):
         # Force eager attention implementation
         use_flash_attention = False  # Override to force eager implementation
         
-        # Check if we're using DeepSpeed
-        using_deepspeed = ds_config_path is not None
+        # Initialize ds_config_path to None before checking
+        ds_config_path = None
+        
+        # Check if DeepSpeed config is available
+        deepspeed_config = config.get("deepspeed_config", None)
+        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
+                
+            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
+        else:
+            logger.warning("No DeepSpeed configuration found - continuing without DeepSpeed")
+            using_deepspeed = False
         
         # Initialize model with our safe loading function
         logger.info("Loading pre-quantized model with eager attention")
@@ -617,35 +649,6 @@ 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
-        deepspeed_config = config.get("deepspeed_config", None)
-        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
-                
-            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")
-        else:
-            logger.warning("No DeepSpeed configuration found - continuing without DeepSpeed")
-            ds_config_path = None
-        
         training_args_dict = {
             "output_dir": output_dir,
             "num_train_epochs": training_config.get("num_train_epochs", 3),