diff --git a/sia/llm_engine/qwq_llm_engine.ipynb b/sia/llm_engine/qwq_llm_engine.ipynb index 6d27632..037343c 100644 --- a/sia/llm_engine/qwq_llm_engine.ipynb +++ b/sia/llm_engine/qwq_llm_engine.ipynb @@ -12,7 +12,7 @@ "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n", "Unsloth: Failed to patch Gemma3ForConditionalGeneration.\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": { "application/vnd.jupyter.widget-view+json": { - "model_id": "a99acfd42fa645abae77144125a734b7", + "model_id": "4f302eb9995d47fab2aa6339aa00a8d8", "version_major": 2, "version_minor": 0 }, @@ -109,7 +109,7 @@ "# Load model\n", "model, _returned_tokenizer = FastLanguageModel.from_pretrained(\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", "execution_count": 6, "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": [ { "name": "stderr", @@ -138,7 +185,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 8, "metadata": {}, "outputs": [], "source": [ @@ -147,7 +194,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 9, "metadata": {}, "outputs": [], "source": [ @@ -160,7 +207,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 10, "metadata": {}, "outputs": [], "source": [ @@ -173,7 +220,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 11, "metadata": {}, "outputs": [], "source": [ @@ -185,7 +232,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 12, "metadata": {}, "outputs": [], "source": [ @@ -200,7 +247,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 13, "metadata": {}, "outputs": [], "source": [ @@ -209,7 +256,7 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 14, "metadata": {}, "outputs": [], "source": [ @@ -223,15 +270,14 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 15, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "\n", - "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?<|im_end|>" + "Hi, I'm an AI assistant. I don't have feelings, but I'm here to help you. How can I assist you today?<|im_end|>" ] } ], @@ -242,7 +288,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 16, "metadata": {}, "outputs": [], "source": [ diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py index 8283681..a0e1dd7 100644 --- a/sia/llm_engine/qwq_llm_engine.py +++ b/sia/llm_engine/qwq_llm_engine.py @@ -6,8 +6,6 @@ from threading import Thread from transformers import AutoTokenizer, TextIteratorStreamer, pipeline from typing import Callable, Iterator, Optional from xml_schema_validator import XmlLogitsProcessor -import os -import torch from . import LlmEngine from .. import util @@ -38,10 +36,7 @@ class QwQLlmEngine(LlmEngine): # Load model self._model, _returned_tokenizer = FastLanguageModel.from_pretrained( model_path, - load_in_4bit = True, # False for LoRA 16bit - fast_inference = True, # Enable vLLM fast inference - gpu_memory_utilization = 0.8, # Reduce if out of memory - tokenizer = self._tokenizer, + gpu_memory_utilization = 0.5, # Reduce if out of memory ) # enable unsloth optimizations diff --git a/tools/train/train/qwq.py b/tools/train/train/qwq.py index d42ed05..a6eea2e 100644 --- a/tools/train/train/qwq.py +++ b/tools/train/train/qwq.py @@ -30,7 +30,7 @@ class Args: parser.add_argument( '--base-model', type=str, - default='Qwen/QwQ-32B', + default='unsloth/QwQ-32B-bnb-4bit', help='HuggingFace model ID for base model' ) parser.add_argument( @@ -71,9 +71,6 @@ def main(): dataset = Dataset(args.config_path) 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: tokenizer_config = json.load(f) @@ -84,32 +81,23 @@ def main(): model, _returned_tokenizer = FastLanguageModel.from_pretrained( model_name = args.base_model, - max_seq_length = max_seq_length, - 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, + gpu_memory_utilization = 0.5, ) model = FastLanguageModel.get_peft_model( model, - r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128 target_modules = [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], # Remove QKVO if out of memory - lora_alpha = lora_rank, lora_dropout = 0, # Supports any, but = 0 is optimized bias = "none", # Supports any, but = "none" is optimized use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context random_state = 3407, - use_rslora = False, # We support rank stabilized LoRA - loftq_config = None, # And LoftQ ) training_args = TrainingArguments( - output_dir=str(args.output_dir), + output_dir=str(args.output_dir) + "_train", num_train_epochs=3, per_device_train_batch_size=1, gradient_accumulation_steps=16, @@ -133,7 +121,6 @@ def main(): args=training_args, train_dataset=dataset.to_transformers_dataset(tokenizer), dataset_text_field="messages", - max_seq_length=max_seq_length, ) trainer.train() @@ -141,6 +128,7 @@ def main(): model.save_pretrained_merged( str(args.output_dir), tokenizer=tokenizer, + save_method="merged_4bit_forced" ) if __name__ == "__main__":