Create README.md

README.md

@@ -0,0 +1,167 @@
+---
+license: mit
+base_model: Qwen/Qwen2.5-Coder-3B
+datasets:
+- GPUMODE/KernelBook
+tags:
+- qwen2
+- code-generation
+- triton
+- pytorch
+- kernel-generation
+- kernelbook
+- lora
+- finetune
+---
+
+# Qwen2.5-Coder-3B-KernelBook: Fine-tuned for PyTorch to Triton Kernel Generation
+
+This repository contains a fine-tuned version of the **[Qwen/Qwen2.5-Coder-3B](https://huggingface.co/Qwen/Qwen2.5-Coder-3B)** model, specialized for transpiling PyTorch `nn.Module` code into high-performance Triton kernels.
+
+The model was trained on the **[GPUMODE/KernelBook](https://huggingface.co/datasets/GPUMODE/KernelBook)** dataset, which contains thousands of pairs of equivalent PyTorch and Triton code snippets generated by `torch.compile`. This fine-tuning enables the model to understand the patterns of PyTorch operations and translate them into efficient, fused GPU kernels written in the Triton language.
+
+This model was fine-tuned as part of a demonstration of an end-to-end workflow: from dataset preparation and model training to benchmarking with the official `KernelBench` framework.
+
+## Model Details
+
+- **Base Model:** `Qwen/Qwen2.5-Coder-3B`
+- **Fine-tuning Dataset:** `GPUMODE/KernelBook`
+- **Method:** Low-Rank Adaptation (LoRA)
+- **Framework:** PyTorch 2.5.0, Transformers, PEFT, TRL
+
+### Training Summary
+
+The model was trained for **1 full epoch** on the `GPUMODE/KernelBook` dataset (18,162 examples), showcasing strong learning and convergence.
+
+- **Final Training Loss:** **`0.0922`**
+- **Final Mean Token Accuracy:** **`98.34%`**
+- **Training Runtime:** `5818.25 seconds` (approx. 1 hour 37 minutes)
+- **Hardware:** 1x NVIDIA H100 80GB
+
+**Key Training Hyperparameters:**
+- `learning_rate`: 2e-4
+- `per_device_train_batch_size`: 1
+- `gradient_accumulation_steps`: 8 (effective batch size of 8)
+- `max_seq_length`: 4096
+- `optimizer`: adamw_torch_fused
+- `precision`: bfloat16
+
+For a detailed view of the training progress, you can visit the [Weights & Biases run page](https://wandb.ai/tarunreddi-university-at-buffalo/huggingface/runs/ew21hn3w).
+
+## How to Use
+
+This model is designed to be used for code generation in a structured prompt format. You should provide the PyTorch code and ask for the Triton code in return.
+
+### Installation
+
+First, make sure you have the necessary libraries installed:
+
+```bash
+pip install torch transformers peft accelerate
+```
+
+### Example Usage
+
+Here is a Python snippet demonstrating how to generate a Triton kernel from a PyTorch `nn.Module`.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# The repository ID of this model on the Hugging Face Hub
+model_id = "TEEN-D/Qwen2.5-Coder-3B-KernelBook-Finetuned"
+
+print("Loading model and tokenizer...")
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True
+)
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+print("Model loaded successfully.")
+
+# --- 1. Define your PyTorch code ---
+pytorch_code = """
+import torch
+import torch.nn as nn
+
+class SumAggregator(nn.Module):
+    def __init__(self):
+        super(SumAggregator, self).__init__()
+
+    def forward(self, neighbor):
+        return torch.sum(neighbor, dim=1)
+"""
+
+# --- 2. Format the prompt as used during training ---
+prompt = f"""### INSTRUCTION
+Generate the Triton code for the following Python code.
+
+### PYTHON CODE:
+{pytorch_code}
+
+### TRITON CODE:
+"""
+
+# --- 3. Generate the Triton kernel ---
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+outputs = model.generate(
+    **inputs,
+    max_new_tokens=2048,
+    do_sample=False, # Use greedy decoding for reproducibility
+    pad_token_id=tokenizer.eos_token_id
+)
+
+full_output = tokenizer.decode(outputs, skip_special_tokens=True)
+
+# --- 4. Extract and print only the Triton code ---
+try:
+    triton_code = full_output.split("### TRITON CODE:").strip()
+    print("\n--- Generated Triton Code ---")
+    print(triton_code)
+except IndexError:
+    print("Could not parse the output. Full generated text:")
+    print(full_output)
+
+```
+
+## Fine-tuning Dataset: GPUMODE/KernelBook
+
+This model's capabilities are a direct result of the high-quality `GPUMODE/KernelBook` dataset.
+
+- **Content:** The dataset contains 18,162 pairs of PyTorch programs and their equivalent Triton kernels, as generated by `torch.compile`.
+- **Creation Process:** The authors collected PyTorch repositories, extracted `nn.Module` classes, generated Triton code with `torch.compile`, and enriched the data with metadata.
+- **Recommended Usage:** For best results when using or evaluating the generated Triton code, it is recommended to use the same PyTorch version the dataset was created with (`torch==2.5.0`).
+
+## Base Model: Qwen2.5-Coder-3B
+
+`Qwen2.5-Coder` is a series of code-specific large language models. The 3B model has the following characteristics:
+- **Parameters:** 3.09B
+- **Context Length:** 32,768 tokens
+- **Architecture:** Transformer with RoPE, SwiGLU, RMSNorm.
+
+## Citation
+
+If you use this model or the dataset in your work, please cite the original authors.
+
+**To cite the dataset:**
+```bibtex
+@software{kernelbook2025,
+    title={KernelBook},
+    author={Paliskara, Sahan and Saroufim, Mark},
+    year={2025},
+    month={5},
+    url={https://huggingface.co/datasets/GPUMODE/KernelBook},
+}
+```
+
+**To cite the base model:**
+```bibtex
+@article{hui2024qwen2,
+      title={Qwen2. 5-Coder Technical Report},
+      author={Hui, Binyuan and Yang, Jian and Cui, Zeyu and Yang, Jiaxi and Liu, Dayiheng and Zhang, Lei and Liu, Tianyu and Zhang, Jiajun and Yu, Bowen and Dang, Kai and others},
+      journal={arXiv preprint arXiv:2409.12186},
+      year={2024}
+}
+```