Update app.py

app.py

@@ -1,17 +1,26 @@
 import os
 import torch
+import torch.distributed as dist
 from unsloth import FastLanguageModel, is_bfloat16_supported
 from trl import SFTTrainer
 from transformers import TrainingArguments
 from datasets import load_dataset
 import gradio as gr
 
-
 max_seq_length = 4096
 dtype = None
 load_in_4bit = True
 hf_token = os.getenv("Token")
 
+def setup_distributed_training():
+    dist.init_process_group(backend='nccl')
+    torch.cuda.set_device(dist.get_rank())
+
+def cleanup_distributed_training():
+    dist.destroy_process_group()
+
+setup_distributed_training()
+
 print("Starting model and tokenizer loading...")
 
 # Load the model and tokenizer
@@ -121,8 +130,10 @@ trainer = SFTTrainer(
         weight_decay=0.01,
         lr_scheduler_type="linear",
         seed=3407,
-        local_rank=4,
         output_dir="outputs",
+        # Distributed training arguments
+        deepspeed=None,  # If using deepspeed for further optimizations
+        local_rank=dist.get_rank(),  # Add this line
     ),
 )
 print("Trainer initialized.")
@@ -142,4 +153,6 @@ model.push_to_hub_merged(
     save_method="merged_16bit",
     token=True
 )
-print("Model pushed to hub successfully.")
+print("Model pushed to hub successfully.")
+
+cleanup_distributed_training()