llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implement 4-bit quant for train script and inference

Browse Source
This commit is contained in:
Niels Geens
2025-04-23 16:43:55 +00:00
parent 2c4b254bb0
commit e2e346d134
3 changed files with 65 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
sia/llm_engine/qwq_llm_engine.ipynb
View File

@@ -12,7 +12,7 @@
"🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n", "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
"Unsloth: Failed to patch Gemma3ForConditionalGeneration.\n", "Unsloth: Failed to patch Gemma3ForConditionalGeneration.\n",
"🦥 Unsloth Zoo will now patch everything to make training faster!\n", "🦥 Unsloth Zoo will now patch everything to make training faster!\n",
"INFO 04-23 16:36:10 [__init__.py:239] Automatically detected platform cuda.\n" "INFO 04-23 16:42:57 [__init__.py:239] Automatically detected platform cuda.\n"
] ]
} }
], ],
@@ -93,7 +93,7 @@
{ {
"data": { "data": {
"application/vnd.jupyter.widget-view+json": { "application/vnd.jupyter.widget-view+json": {
"model_id": "a99acfd42fa645abae77144125a734b7", "model_id": "4f302eb9995d47fab2aa6339aa00a8d8",
"version_major": 2, "version_major": 2,
"version_minor": 0 "version_minor": 0
}, },
@@ -109,7 +109,7 @@
"# Load model\n", "# Load model\n",
"model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n", "model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
" model_path,\n", " model_path,\n",
" gpu_memory_utilization = 0.8, # Reduce if out of memory\n", " gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
")" ")"
] ]
}, },
@@ -117,6 +117,53 @@
"cell_type": "code", "cell_type": "code",
"execution_count": 6, "execution_count": 6,
"metadata": {}, "metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Qwen2ForCausalLM(\n",
" (model): Qwen2Model(\n",
" (embed_tokens): Embedding(152064, 5120, padding_idx=151654)\n",
" (layers): ModuleList(\n",
" (0-63): 64 x Qwen2DecoderLayer(\n",
" (self_attn): Qwen2Attention(\n",
" (q_proj): Linear4bit(in_features=5120, out_features=5120, bias=True)\n",
" (k_proj): Linear4bit(in_features=5120, out_features=1024, bias=True)\n",
" (v_proj): Linear4bit(in_features=5120, out_features=1024, bias=True)\n",
" (o_proj): Linear4bit(in_features=5120, out_features=5120, bias=False)\n",
" (rotary_emb): LlamaRotaryEmbedding()\n",
" )\n",
" (mlp): Qwen2MLP(\n",
" (gate_proj): Linear4bit(in_features=5120, out_features=27648, bias=False)\n",
" (up_proj): Linear4bit(in_features=5120, out_features=27648, bias=False)\n",
" (down_proj): Linear4bit(in_features=27648, out_features=5120, bias=False)\n",
" (act_fn): SiLU()\n",
" )\n",
" (input_layernorm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
" (post_attention_layernorm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
" )\n",
" )\n",
" (norm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
" (rotary_emb): LlamaRotaryEmbedding()\n",
" )\n",
" (lm_head): Linear(in_features=5120, out_features=152064, bias=False)\n",
")"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# enable unsloth optimizations\n",
"FastLanguageModel.for_inference(model)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [ "outputs": [
{ {
"name": "stderr", "name": "stderr",
@@ -138,7 +185,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 7, "execution_count": 8,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -147,7 +194,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 8, "execution_count": 9,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -160,7 +207,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 9, "execution_count": 10,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -173,7 +220,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 10, "execution_count": 11,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -185,7 +232,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 11, "execution_count": 12,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -200,7 +247,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 12, "execution_count": 13,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -209,7 +256,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 13, "execution_count": 14,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -223,15 +270,14 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 14, "execution_count": 15,
"metadata": {}, "metadata": {},
"outputs": [ "outputs": [
{ {
"name": "stdout", "name": "stdout",
"output_type": "stream", "output_type": "stream",
"text": [ "text": [
"<write_stdout>\n", "<write_stdout>Hi, I'm an AI assistant. I don't have feelings, but I'm here to help you. How can I assist you today?</write_stdout><|im_end|>"
"Hello! I'm just a large language model, so I don't have feelings, but thank you for asking. How can I assist you today?</write_stdout><|im_end|>"
] ]
} }
], ],
@@ -242,7 +288,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 15, "execution_count": 16,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [

7
sia/llm_engine/qwq_llm_engine.py
View File

@@ -6,8 +6,6 @@ from threading import Thread
from transformers import AutoTokenizer, TextIteratorStreamer, pipeline from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
from typing import Callable, Iterator, Optional from typing import Callable, Iterator, Optional
from xml_schema_validator import XmlLogitsProcessor from xml_schema_validator import XmlLogitsProcessor
import os
import torch
from . import LlmEngine from . import LlmEngine
from .. import util from .. import util
@@ -38,10 +36,7 @@ class QwQLlmEngine(LlmEngine):
# Load model # Load model
self._model, _returned_tokenizer = FastLanguageModel.from_pretrained( self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
model_path, model_path,
load_in_4bit = True, # False for LoRA 16bit gpu_memory_utilization = 0.5, # Reduce if out of memory
fast_inference = True, # Enable vLLM fast inference
gpu_memory_utilization = 0.8, # Reduce if out of memory
tokenizer = self._tokenizer,
) )
# enable unsloth optimizations # enable unsloth optimizations

20
tools/train/train/qwq.py
View File

@@ -30,7 +30,7 @@ class Args:
parser.add_argument( parser.add_argument(
'--base-model', '--base-model',
type=str, type=str,
default='Qwen/QwQ-32B', default='unsloth/QwQ-32B-bnb-4bit',
help='HuggingFace model ID for base model' help='HuggingFace model ID for base model'
) )
parser.add_argument( parser.add_argument(
@@ -71,9 +71,6 @@ def main():
dataset = Dataset(args.config_path) dataset = Dataset(args.config_path)
dataset.validate() dataset.validate()
max_seq_length = 2048 # Can increase for longer reasoning traces
lora_rank = 64 # Larger rank = smarter, but slower
with open('/root/sia/qwq_tokenizer_config.json', 'r') as f: with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
tokenizer_config = json.load(f) tokenizer_config = json.load(f)
@@ -84,32 +81,23 @@ def main():
model, _returned_tokenizer = FastLanguageModel.from_pretrained( model, _returned_tokenizer = FastLanguageModel.from_pretrained(
model_name = args.base_model, model_name = args.base_model,
max_seq_length = max_seq_length, gpu_memory_utilization = 0.5,
load_in_4bit = True, # False for LoRA 16bit
fast_inference = True, # Enable vLLM fast inference
max_lora_rank = lora_rank,
gpu_memory_utilization = 0.5, # Reduce if out of memory
tokenizer = tokenizer,
) )
model = FastLanguageModel.get_peft_model( model = FastLanguageModel.get_peft_model(
model, model,
r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
target_modules = [ target_modules = [
"q_proj", "k_proj", "v_proj", "o_proj", "q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj", "gate_proj", "up_proj", "down_proj",
], # Remove QKVO if out of memory ], # Remove QKVO if out of memory
lora_alpha = lora_rank,
lora_dropout = 0, # Supports any, but = 0 is optimized lora_dropout = 0, # Supports any, but = 0 is optimized
bias = "none", # Supports any, but = "none" is optimized bias = "none", # Supports any, but = "none" is optimized
use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
random_state = 3407, random_state = 3407,
use_rslora = False, # We support rank stabilized LoRA
loftq_config = None, # And LoftQ
) )
training_args = TrainingArguments( training_args = TrainingArguments(
output_dir=str(args.output_dir), output_dir=str(args.output_dir) + "_train",
num_train_epochs=3, num_train_epochs=3,
per_device_train_batch_size=1, per_device_train_batch_size=1,
gradient_accumulation_steps=16, gradient_accumulation_steps=16,
@@ -133,7 +121,6 @@ def main():
args=training_args, args=training_args,
train_dataset=dataset.to_transformers_dataset(tokenizer), train_dataset=dataset.to_transformers_dataset(tokenizer),
dataset_text_field="messages", dataset_text_field="messages",
max_seq_length=max_seq_length,
) )
trainer.train() trainer.train()
@@ -141,6 +128,7 @@ def main():
model.save_pretrained_merged( model.save_pretrained_merged(
str(args.output_dir), str(args.output_dir),
tokenizer=tokenizer, tokenizer=tokenizer,
save_method="merged_4bit_forced"
) )
if __name__ == "__main__": if __name__ == "__main__":