Upload 12 files

app.py

@@ -0,0 +1,212 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import torch
+import random
+import time
+import os
+
+# Load the model and tokenizer
+model_path = "./phi2-qlora-final"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
+
+# Custom CSS for better styling
+custom_css = """
+.gradio-container {
+    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+}
+.container {
+    max-width: 800px;
+    margin: auto;
+    padding: 20px;
+}
+.title {
+    text-align: center;
+    color: #2c3e50;
+    margin-bottom: 20px;
+}
+.description {
+    text-align: center;
+    color: #7f8c8d;
+    margin-bottom: 30px;
+}
+.loading {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    height: 100px;
+}
+.error {
+    color: #e74c3c;
+    padding: 10px;
+    border-radius: 5px;
+    background-color: #fde8e8;
+    margin: 10px 0;
+}
+"""
+
+def generate_response(prompt, max_length=512, temperature=0.7, top_p=0.9, top_k=50):
+    """Generate response with progress indicator"""
+    try:
+        if not prompt.strip():
+            return "Please enter a prompt."
+            
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        outputs = model.generate(
+            **inputs,
+            max_length=max_length,
+            temperature=temperature,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id,
+            do_sample=True,
+            top_p=top_p,
+            top_k=top_k,
+        )
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
+
+def clear_all():
+    """Clear all inputs and outputs"""
+    return "", "", 512, 0.7, 0.9, 50
+
+# Example prompts
+example_prompts = [
+    "What is the capital of France?",
+    "Explain quantum computing in simple terms.",
+    "Write a short story about a robot learning to paint.",
+    "What are the benefits of meditation?",
+    "How does photosynthesis work?",
+]
+
+# Create the Gradio interface
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as iface:
+    gr.Markdown(
+        """
+        # 🤖 Phi-2 QLoRA Chat Interface
+        
+        Chat with the fine-tuned Phi-2 model using QLoRA. Adjust the parameters below to control the generation.
+        """,
+        elem_classes="title"
+    )
+    
+    gr.Markdown(
+        """
+        This interface allows you to interact with a fine-tuned Phi-2 model. You can adjust various parameters to control the generation process.
+        """,
+        elem_classes="description"
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            # Input section
+            with gr.Group():
+                gr.Markdown("### 💭 Input")
+                prompt = gr.Textbox(
+                    label="Enter your prompt:",
+                    placeholder="Type your message here...",
+                    lines=3,
+                    show_label=True,
+                    container=True
+                )
+                
+                with gr.Row():
+                    max_length = gr.Slider(
+                        minimum=64,
+                        maximum=1024,
+                        value=512,
+                        step=64,
+                        label="Max Length",
+                        info="Maximum length of generated response"
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.7,
+                        step=0.1,
+                        label="Temperature",
+                        info="Higher values make output more random"
+                    )
+                
+                with gr.Row():
+                    top_p = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.9,
+                        step=0.1,
+                        label="Top P",
+                        info="Nucleus sampling parameter"
+                    )
+                    top_k = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        value=50,
+                        step=1,
+                        label="Top K",
+                        info="Top-k sampling parameter"
+                    )
+            
+            # Buttons
+            with gr.Row():
+                submit_btn = gr.Button("Generate Response", variant="primary")
+                clear_btn = gr.Button("Clear All", variant="secondary")
+        
+        with gr.Column(scale=2):
+            # Output section
+            with gr.Group():
+                gr.Markdown("### 🤖 Response")
+                output = gr.Textbox(
+                    label="Model Response:",
+                    lines=5,
+                    show_label=True,
+                    container=True
+                )
+    
+    # Examples section
+    with gr.Group():
+        gr.Markdown("### 📝 Example Prompts")
+        gr.Examples(
+            examples=example_prompts,
+            inputs=prompt,
+            outputs=output,
+            fn=generate_response,
+            cache_examples=True
+        )
+    
+    # Footer
+    gr.Markdown(
+        """
+        ---
+        Made with ❤️ using Phi-2 and QLoRA
+        """,
+        elem_classes="footer"
+    )
+    
+    # Event handlers
+    submit_btn.click(
+        fn=generate_response,
+        inputs=[prompt, max_length, temperature, top_p, top_k],
+        outputs=output
+    )
+    
+    clear_btn.click(
+        fn=clear_all,
+        inputs=[],
+        outputs=[prompt, output, max_length, temperature, top_p, top_k]
+    )
+
+if __name__ == "__main__":
+    iface.launch(
+        share=True,  # Enable sharing
+        server_name="0.0.0.0",  # Allow external access
+        server_port=7860,  # Default Gradio port
+        show_error=True  # Show detailed error messages
+    ) 

phi2-qlora-final/README.md

@@ -0,0 +1,202 @@
+---
+base_model: microsoft/phi-2
+library_name: peft
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.14.0

phi2-qlora-final/adapter_config.json

@@ -0,0 +1,34 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "microsoft/phi-2",
+  "bias": "none",
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 8,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "o_proj",
+    "q_proj",
+    "v_proj",
+    "k_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

phi2-qlora-final/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc69559ce651a6a9261ad792f1edaf98f436d54e52df0ce004b57d88672b2367
+size 15754072

phi2-qlora-final/added_tokens.json

@@ -0,0 +1,40 @@
+{
+  "\t\t": 50294,
+  "\t\t\t": 50293,
+  "\t\t\t\t": 50292,
+  "\t\t\t\t\t": 50291,
+  "\t\t\t\t\t\t": 50290,
+  "\t\t\t\t\t\t\t": 50289,
+  "\t\t\t\t\t\t\t\t": 50288,
+  "\t\t\t\t\t\t\t\t\t": 50287,
+  "  ": 50286,
+  "   ": 50285,
+  "    ": 50284,
+  "     ": 50283,
+  "      ": 50282,
+  "       ": 50281,
+  "        ": 50280,
+  "         ": 50279,
+  "          ": 50278,
+  "           ": 50277,
+  "            ": 50276,
+  "             ": 50275,
+  "              ": 50274,
+  "               ": 50273,
+  "                ": 50272,
+  "                 ": 50271,
+  "                  ": 50270,
+  "                   ": 50269,
+  "                    ": 50268,
+  "                     ": 50267,
+  "                      ": 50266,
+  "                       ": 50265,
+  "                        ": 50264,
+  "                         ": 50263,
+  "                          ": 50262,
+  "                           ": 50261,
+  "                            ": 50260,
+  "                             ": 50259,
+  "                              ": 50258,
+  "                               ": 50257
+}

phi2-qlora-final/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

phi2-qlora-final/tokenizer_config.json

@@ -0,0 +1,326 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50257": {
+      "content": "                               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "                         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50277": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50280": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50281": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50282": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50283": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50284": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50285": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50286": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50287": {
+      "content": "\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50288": {
+      "content": "\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50289": {
+      "content": "\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50290": {
+      "content": "\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50291": {
+      "content": "\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50292": {
+      "content": "\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50293": {
+      "content": "\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50294": {
+      "content": "\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 2048,
+  "pad_token": "<|endoftext|>",
+  "return_token_type_ids": false,
+  "tokenizer_class": "CodeGenTokenizer",
+  "unk_token": "<|endoftext|>"
+}

phi2-qlora-final/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:080cb7f5521076a88a3e9b163a231d1fb82959fccc6ccbaa0e170c6a6c1eb6c5
+size 5560

train.py

@@ -0,0 +1,176 @@
+import os
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    TrainingArguments,
+    pipeline,
+    logging,
+    DataCollatorForLanguageModeling,
+)
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from trl import SFTTrainer
+import torch
+import logging
+from torch.utils.data import DataLoader
+import multiprocessing
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def preprocess_function(examples, tokenizer):
+    # Format the text
+    texts = []
+    for i in range(len(examples["text"])):
+        text = examples["text"][i]
+        texts.append(text)
+    
+    # Tokenize the texts with shorter max length
+    tokenized = tokenizer(
+        texts,
+        padding=True,
+        truncation=True,
+        max_length=512,  # Reduced from 1024 to 512
+        return_tensors="pt"
+    )
+    
+    return tokenized
+
+def main():
+    try:
+        # Load dataset
+        logger.info("Loading dataset...")
+        dataset = load_dataset("OpenAssistant/oasst1")
+        
+        # Use a smaller subset for faster training
+        logger.info("Selecting smaller dataset subset...")
+        dataset["train"] = dataset["train"].select(range(2000))  # Reduced to 2k examples
+        
+        # Model and tokenizer setup
+        logger.info("Setting up model and tokenizer...")
+        model_name = "microsoft/phi-2"
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer.pad_token = tokenizer.eos_token
+
+        # Preprocess dataset
+        logger.info("Preprocessing dataset...")
+        tokenized_dataset = dataset.map(
+            lambda x: preprocess_function(x, tokenizer),
+            batched=True,
+            remove_columns=dataset["train"].column_names,
+            num_proc=4  # Parallel processing for faster preprocessing
+        )
+        
+        # Split dataset into train and eval
+        logger.info("Splitting dataset into train and eval sets...")
+        split_dataset = tokenized_dataset["train"].train_test_split(test_size=0.1, seed=42)
+        train_dataset = split_dataset["train"]
+        eval_dataset = split_dataset["test"]
+
+        # Configure 4-bit quantization with memory optimizations
+        logger.info("Configuring 4-bit quantization...")
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_storage=torch.float16
+        )
+
+        # Load model with quantization and memory optimizations
+        logger.info("Loading model...")
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True,
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True
+        )
+
+        # Enable gradient checkpointing for memory efficiency
+        model.gradient_checkpointing_enable()
+        model.enable_input_require_grads()
+
+        # Prepare model for k-bit training
+        logger.info("Preparing model for k-bit training...")
+        model = prepare_model_for_kbit_training(model)
+
+        # LoRA configuration with optimized parameters
+        logger.info("Configuring LoRA...")
+        lora_config = LoraConfig(
+            r=8,  # Reduced from 16 to 8
+            lora_alpha=16,  # Reduced from 32 to 16
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM"
+        )
+
+        # Get PEFT model
+        logger.info("Getting PEFT model...")
+        model = get_peft_model(model, lora_config)
+
+        # Create data collator
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False
+        )
+
+        # Training arguments with memory-optimized settings
+        logger.info("Setting up training arguments...")
+        training_args = TrainingArguments(
+            output_dir="./phi2-qlora",
+            num_train_epochs=2,
+            per_device_train_batch_size=4,  # Reduced from 16 to 4
+            per_device_eval_batch_size=4,
+            gradient_accumulation_steps=4,  # Increased from 1 to 4
+            learning_rate=2e-4,
+            fp16=True,
+            logging_steps=5,
+            save_strategy="epoch",
+            evaluation_strategy="epoch",
+            # Additional optimizations
+            dataloader_num_workers=2,  # Reduced from 4 to 2
+            dataloader_pin_memory=True,
+            warmup_ratio=0.05,
+            lr_scheduler_type="cosine",
+            optim="adamw_torch",
+            max_grad_norm=1.0,
+            group_by_length=True,
+        )
+
+        # Create trainer
+        logger.info("Creating trainer...")
+        trainer = SFTTrainer(
+            model=model,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            tokenizer=tokenizer,
+            args=training_args,
+            data_collator=data_collator,
+        )
+
+        # Train the model
+        logger.info("Starting training...")
+        trainer.train()
+
+        # Save the model
+        logger.info("Saving model...")
+        trainer.save_model("./phi2-qlora-final")
+
+        # Save tokenizer
+        logger.info("Saving tokenizer...")
+        tokenizer.save_pretrained("./phi2-qlora-final")
+        
+        logger.info("Training completed successfully!")
+
+    except Exception as e:
+        logger.error(f"An error occurred: {str(e)}")
+        raise
+
+if __name__ == "__main__":
+    multiprocessing.set_start_method('spawn')
+    main()