aijadugar/Qwen2-0.5B-full-instruct

Qwen2 0.5B Full Fine Tuning for Instruction Following

Overview

This project contains a full fine tuning workflow for the Qwen2 0.5B language model with the objective of improving instruction following capabilities.

The model is trained using supervised fine tuning on an instruction based dataset. 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 fine tuning

Training objective: Instruction following

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 rather than applying LoRA or another parameter efficient fine tuning method.

The training pipeline includes:

1. Loading the pretrained Qwen2 0.5B model
2. Preparing the instruction following dataset
3. Tokenizing the training examples
4. Training the complete model
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 optimizer therefore follows the structure:

BF16 parameters

BF16 gradients

FP32 optimizer states

FP32 AdamW computation

BF16 updated parameters

The CUDA implementation was validated against reference AdamW calculations before being used for model training.

Precision

The training workflow uses BF16 model parameters.

BF16 is used for model parameters and gradients while FP32 is retained for the Adam optimizer states.

This provides a better numerical representation for optimizer statistics than storing the Adam states directly in BF16.

Repository Structure

ft/
├── README.md
├── model.safetensors
├── config.json
├── tokenizer.json
├── tokenizer_config.json
├── special_tokens_map.json
└── other tokenizer and model files

The exact files may vary depending on the Hugging Face Transformers and tokenizer versions used during training.

Loading the Fine Tuned Model

The fine tuned model can be loaded using the Hugging Face Transformers library.

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()

Running Inference

A simple instruction can be provided to the fine tuned model as follows:

prompt = "Explain what machine learning is in simple terms."

inputs = tokenizer(
    prompt,
    return_tensors="pt"
)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=128
    )

response = tokenizer.decode(
    outputs[0],
    skip_special_tokens=True
)

print(response)

For chat formatted datasets, the appropriate Qwen chat template should be applied before generation.

Training Environment

The training workflow was developed and tested in a GPU notebook environment.

Example hardware:

GPU: NVIDIA Tesla T4
GPU configuration: 2 GPUs
Model: Qwen2 0.5B
Training type: Full fine tuning
Precision: BF16

Actual training performance depends on GPU availability, batch size, sequence length, gradient accumulation, dataset size, and the selected Transformers and PyTorch versions.

Requirements

Install the main dependencies with:

pip install torch
pip install transformers
pip install datasets
pip install accelerate
pip install safetensors

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

Output

The final model is stored in Safetensors format.

The primary model file is:

model.safetensors

The model configuration and tokenizer files are stored alongside the weights so that the resulting directory can be loaded directly using the Hugging Face Transformers API.

Notes

This project performs full fine tuning rather than parameter efficient fine tuning.

Because all model parameters are updated, memory requirements are significantly higher than LoRA or QLoRA based training.

The custom CUDA optimizer is intended as an optimization and experimentation component of the training workflow. Its correctness should be validated against a reference AdamW implementation before performing long training runs.

License

The licensing terms of the original Qwen2 model and the training dataset apply to the corresponding components used in this project.

Refer to the original model and dataset licenses before redistributing the fine tuned model or dataset.

Acknowledgements

This project uses the Qwen2 model architecture and the Hugging Face Transformers ecosystem.

Original model documentation and licensing should be consulted before using the model in production or redistributing derived model weights.