aijadugar/Qwen2-0.5B-sft-instruct

Qwen2 0.5B SFT (Full Parameter) for Instruction Following

Overview

This project contains a supervised fine-tuning (SFT) workflow for the Qwen2 0.5B language model with the objective of improving instruction following capabilities.

The model is trained using full-parameter supervised fine tuning on an instruction based dataset, with the loss masked to the assistant response tokens only. The project also includes a custom CUDA fused AdamW optimizer designed to reduce the number of CUDA kernel launches during optimizer updates.

Model

Base model: Qwen2 0.5B

Training method: Full parameter supervised fine tuning

Training objective: Instruction following (prompt-masked loss)

Parameter efficient fine tuning: No

Precision: BF16

Optimizer: Custom fused AdamW CUDA optimizer

Model format: Safetensors

Training

The model is trained by updating all trainable parameters. The training pipeline includes:

1. Loading the pretrained Qwen2 0.5B model
2. Preparing the instruction following dataset
3. Tokenizing the training examples and masking the prompt portion of each label sequence
4. Training the complete model with loss computed only on response tokens
5. Using gradient accumulation
6. Using a linear learning rate schedule with warmup
7. Using BF16 computation
8. Using a custom CUDA fused AdamW optimizer
9. Saving the final model in Safetensors format

Custom CUDA Optimizer

This project includes a custom fused AdamW CUDA implementation. The optimizer performs the following operations inside a CUDA kernel:

1. Adam first moment update
2. Adam second moment update
3. Bias correction
4. Decoupled weight decay
5. Parameter update

The model parameters and gradients use BF16 while the Adam optimizer states and optimizer arithmetic use FP32. The CUDA implementation was validated against reference AdamW calculations before being used for model training.

Loading the Fine Tuned Model

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_path = "./ft"

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16)
model.eval()

Requirements

pip install torch transformers datasets accelerate safetensors

CUDA and a compatible NVIDIA driver are required for CUDA based training.