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:
Loading the pretrained Qwen2 0.5B model
Preparing the instruction following dataset
Tokenizing the training examples and masking the prompt portion of each label sequence
Training the complete model with loss computed only on response tokens
Using gradient accumulation
Using a linear learning rate schedule with warmup
Using BF16 computation
Using a custom CUDA fused AdamW optimizer
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:
Adam first moment update
Adam second moment update
Bias correction
Decoupled weight decay
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.
