Update app.py

app.py

@@ -8,6 +8,9 @@ import gradio as gr
 import json
 from huggingface_hub import HfApi
 
+# Ensure that all 4 GPUs are visible to PyTorch
+os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
+
 max_seq_length = 4096
 dtype = None
 load_in_4bit = True
@@ -17,7 +20,6 @@ current_num = os.getenv("NUM")
 print(f"stage ${current_num}")
 
 api = HfApi(token=hf_token)
-# models = f"dad1909/cybersentinal-2.0-{current_num}"
 
 model_base = "unsloth/llama-3-8b-Instruct-bnb-4bit"
 
@@ -31,13 +33,17 @@ model, tokenizer = FastLanguageModel.from_pretrained(
     load_in_4bit=load_in_4bit,
     token=hf_token
 )
-print("Model and tokenizer loaded successfully.")
+
+# Move the model to GPU
+model = model.to('cuda')
 
 # Wrap the model in DataParallel to use all GPUs
 if torch.cuda.device_count() > 1:
     print(f"Using {torch.cuda.device_count()} GPUs!")
     model = torch.nn.DataParallel(model)
 
+print("Model and tokenizer loaded successfully.")
+
 print("Configuring PEFT model...")
 model = FastLanguageModel.get_peft_model(
     model.module if isinstance(model, torch.nn.DataParallel) else model,
@@ -118,7 +124,7 @@ trainer = SFTTrainer(
     dataset_num_proc=2,
     packing=False,
     args=TrainingArguments(
-        per_device_train_batch_size=17,
+        per_device_train_batch_size=17,  # Adjust this based on GPU memory
         gradient_accumulation_steps=17,
         learning_rate=2e-4,
         fp16=not is_bfloat16_supported(),
@@ -144,21 +150,27 @@ num += 1
 uploads_models = f"cybersentinal-3.0"
 
 print("Saving the trained model...")
-model.module.save_pretrained_merged("model", tokenizer, save_method="merged_16bit") if isinstance(model, torch.nn.DataParallel) else model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")
+if isinstance(model, torch.nn.DataParallel):
+    model.module.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")
+else:
+    model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")
 print("Model saved successfully.")
 
 print("Pushing the model to the hub...")
-model.module.push_to_hub_merged(
-    uploads_models,
-    tokenizer,
-    save_method="merged_16bit",
-    token=hf_token
-) if isinstance(model, torch.nn.DataParallel) else model.push_to_hub_merged(
-    uploads_models,
-    tokenizer,
-    save_method="merged_16bit",
-    token=hf_token
-)
+if isinstance(model, torch.nn.DataParallel):
+    model.module.push_to_hub_merged(
+        uploads_models,
+        tokenizer,
+        save_method="merged_16bit",
+        token=hf_token
+    )
+else:
+    model.push_to_hub_merged(
+        uploads_models,
+        tokenizer,
+        save_method="merged_16bit",
+        token=hf_token
+    )
 print("Model pushed to hub successfully.")
 
 api.delete_space_variable(repo_id="dad1909/CyberCode", key="NUM")