low-precision-cuda

# Dataset Card for CUDA-Mixed-Precision-SFT ## Dataset Summary This dataset contains **1,955** pairs of (FP32 CUDA Kernel -> Optimized Mixed-Precision Kernel) derived from the `mixed_precision_optimizer` project. It is designed to train Large Language Models (LLMs) to automatically optimize CUDA code by leveraging: - FP16/BF16 data types for storage and computation. - Tensor Core (`wmma` / `mma`) instructions where applicable. - Vectorized memory access (`float4`, `half2`). - Stability-aware casting (keeping reductions in FP32). ## Dataset Structure The dataset is provided in JSONL format with the following subsets: | File | Count | Description | |------|-------|-------------| | `full_dataset_sft.jsonl` | 1,955 | Complete dataset including all successful optimizations (correctness verified). **Recommended for Base SFT.** | | `gold_dataset_sft.jsonl` | 732 | High-quality subset where **Speedup > 1.0x** (vs Torch Native). Use for weighted training or preference ranking. | | `train_sft.jsonl` | 1,564 | 80% Training Split | | `val_sft.jsonl` | 195 | 10% Validation Split | ## Data Fields Each sample contains: - `instruction`: The system prompt/instruction given to the model. - `input`: The original unoptimized FP32 CUDA kernel code. - `output`: The optimized Mixed-Precision CUDA kernel code (verified correct). - `metadata`: - `speedup_vs_torch_native`: Performance gain ratio. - `gpu_arch`: Target GPU architecture (e.g., sm80). - `task_type`: Operator type (e.g., `Conv2d`, `Gemm`). ## How to Use ### 1. With HuggingFace `datasets` ```python from datasets import load_dataset dataset = load_dataset("json", data_files={ "train": "llm_training_dataset/train_sft.jsonl", "validation": "llm_training_dataset/val_sft.jsonl" }) print(dataset["train"][0]) ``` ### 2. With LLaMA-Factory (Training) Add the following to your `dataset_info.json`: ```json "cuda_mixed_precision": { "file_name": "full_dataset_sft.jsonl", "columns": { "prompt": "instruction", "query": "input", "response": "output" } } ``` ### 3. Prompt Format To get the best results from a model trained on this dataset, use the following prompt template at inference time: ```text 将以下FP32 CUDA kernel优化为混合精度实现。要求： 1. 使用FP16/BF16进行内部计算，保持接口float*不变 2. 关键操作（reduction/softmax/exp）保留FP32累加确保数值稳定性 3. 使用half2/float4向量化优化内存访问 4. 生成完整可编译的.cu文件，包含PYBIND11_MODULE定义 [INPUT_CODE_HERE] ``` ## Considerations for Use - **Architecture Specificity**: Most kernels are optimized for NVIDIA Ampere (sm80) architectures. - **Correctness vs Performance**: The "Full" dataset guarantees correctness but not always performance improvement. For performance-critical tasks, prioritize the "Gold" samples or use Reinforcement Learning (DPO) with the provided speedup metrics.