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

Flash attention in notebook

Browse Source
This commit is contained in:
Niels Geens
2025-04-23 19:38:56 +00:00
parent 4b124fa3ed
commit 9b2c2ece70
1 changed files with 38 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files

170
sia/llm_engine/qwq_llm_engine.ipynb
View File

@@ -1,21 +1,19 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"cell_type": "markdown",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"🦥 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:42:57 [__init__.py:239] Automatically detected platform cuda.\n"
]
}
],
"source": [
"source /root/venvs/sia/bin/activate\n",
"apt-get update && apt-get install -y cuda-toolkit\n",
"pip install flash-attn --no-build-isolation"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Unsloth should be imported before transformers to ensure all optimizations are applied.\n",
"from unsloth import FastLanguageModel"
@@ -23,7 +21,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -35,126 +33,48 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"temperature = 0.6\n",
"model_path = \"/root/models/notebook_merged_4bit\"\n",
"model_path = \"/root/models/current\"\n",
"xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
]
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Load tokenizer\n",
"tokenizer = AutoTokenizer.from_pretrained(\n",
" model_path,\n",
" legacy=False,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"==((====))== Unsloth 2025.3.19: Fast Qwen2 patching. Transformers: 4.51.3. vLLM: 0.8.2.\n",
" \\\\ /| NVIDIA RTX 6000 Ada Generation. Num GPUs = 1. Max memory: 47.5 GB. Platform: Linux.\n",
"O^O/ \\_/ \\ Torch: 2.6.0+cu124. CUDA: 8.9. CUDA Toolkit: 12.4. Triton: 3.2.0\n",
"\\ / Bfloat16 = TRUE. FA [Xformers = 0.0.29.post2. FA2 = False]\n",
" \"-____-\" Free license: http://github.com/unslothai/unsloth\n",
"Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Sliding Window Attention is enabled but not implemented for `eager`; unexpected results may be encountered.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"validate_env\n",
"device_map sequential\n",
"validate_env\n",
"device_map OrderedDict([('', 0)])\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "4f302eb9995d47fab2aa6339aa00a8d8",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading checkpoint shards: 0%| | 0/4 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"outputs": [],
"source": [
"# Load model\n",
"model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
" model_path,\n",
" gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
" load_in_4bit=True,\n",
" attn_implementation=\"flash_attention_2\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"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"
}
],
"outputs": [],
"source": [
"# enable unsloth optimizations\n",
"FastLanguageModel.for_inference(model)"
@@ -162,30 +82,24 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Device set to use cuda:0\n"
]
}
],
"outputs": [],
"source": [
"import torch\n",
"# Create inference pipeline with memory-efficient settings\n",
"pipeline = pipeline(\n",
" \"text-generation\",\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" return_full_text=False,\n",
" torch_dtype=torch.float16,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -194,7 +108,7 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -207,7 +121,7 @@
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -220,7 +134,7 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -232,7 +146,7 @@
},
{
"cell_type": "code",
"execution_count": 12,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -247,7 +161,7 @@
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -256,7 +170,7 @@
},
{
"cell_type": "code",
"execution_count": 14,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -270,17 +184,9 @@
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<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|>"
]
}
],
"outputs": [],
"source": [
"for text in streamer:\n",
" print(text, end=\"\")"
@@ -288,7 +194,7 @@
},
{
"cell_type": "code",
"execution_count": 16,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -301,7 +207,7 @@
"kernelspec": {
"display_name": "sia",
"language": "python",
"name": "python3"
"name": "sia"
},
"language_info": {
"codemirror_mode": {