mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
Optimize GPU memory usage by running inference in inference_mode and clearing caches
This commit is contained in:
Sarvesh
parent
a6d558b44d
commit
2ae46f8908
1 changed files with 149 additions and 122 deletions
271
demo.ipynb
271
demo.ipynb
|
|
@ -567,50 +567,62 @@
|
|||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import torch\n",
|
||||
"import gc\n",
|
||||
"\n",
|
||||
"# Ensure the model is reset to the clean base model state (LoRA hooks disabled)\n",
|
||||
"model.reset()\n",
|
||||
"torch.cuda.empty_cache()\n",
|
||||
"gc.collect()\n",
|
||||
"\n",
|
||||
"base_predictions = []\n",
|
||||
"print(\"Running Base Model Inference (with visual tokens) for all 6 examples...\")\n",
|
||||
"for idx, example in enumerate(examples):\n",
|
||||
" print(f\"Processing Example {idx+1}/{len(examples)}: {example['id']}...\")\n",
|
||||
" \n",
|
||||
" base_messages = [\n",
|
||||
" [\n",
|
||||
" {\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"content\": [\n",
|
||||
" {\"type\": \"video\", \"path\": example[\"video_path\"]},\n",
|
||||
" {\"type\": \"text\", \"text\": example[\"prompt\"]}\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
"with torch.inference_mode():\n",
|
||||
" for idx, example in enumerate(examples):\n",
|
||||
" print(f\"Processing Example {idx+1}/{len(examples)}: {example['id']}...\")\n",
|
||||
" \n",
|
||||
" base_messages = [\n",
|
||||
" [\n",
|
||||
" {\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"content\": [\n",
|
||||
" {\"type\": \"video\", \"path\": example[\"video_path\"]},\n",
|
||||
" {\"type\": \"text\", \"text\": example[\"prompt\"]}\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
" ]\n",
|
||||
" ]\n",
|
||||
" ]\n",
|
||||
"\n",
|
||||
" base_inputs = prepare_smolvlm_inputs(\n",
|
||||
" processor,\n",
|
||||
" base_messages,\n",
|
||||
" device,\n",
|
||||
" video_fps=train_args.video_fps,\n",
|
||||
" max_frames=train_args.max_frames,\n",
|
||||
" video_size_longest_edge=train_args.video_size_longest_edge,\n",
|
||||
" video_load_backend=train_args.video_load_backend\n",
|
||||
" )\n",
|
||||
" base_inputs = prepare_smolvlm_inputs(\n",
|
||||
" processor,\n",
|
||||
" base_messages,\n",
|
||||
" device,\n",
|
||||
" video_fps=train_args.video_fps,\n",
|
||||
" max_frames=train_args.max_frames,\n",
|
||||
" video_size_longest_edge=train_args.video_size_longest_edge,\n",
|
||||
" video_load_backend=train_args.video_load_backend\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" base_generated_ids = raw_model.generate(\n",
|
||||
" **base_inputs,\n",
|
||||
" max_new_tokens=train_args.generation_max_new_tokens,\n",
|
||||
" do_sample=False,\n",
|
||||
" pad_token_id=tokenizer.pad_token_id,\n",
|
||||
" eos_token_id=tokenizer.eos_token_id\n",
|
||||
" )\n",
|
||||
" base_generated_ids = raw_model.generate(\n",
|
||||
" **base_inputs,\n",
|
||||
" max_new_tokens=train_args.generation_max_new_tokens,\n",
|
||||
" do_sample=False,\n",
|
||||
" pad_token_id=tokenizer.pad_token_id,\n",
|
||||
" eos_token_id=tokenizer.eos_token_id\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" base_pred = tokenizer.decode(\n",
|
||||
" base_generated_ids[0][base_inputs[\"input_ids\"].shape[1]:],\n",
|
||||
" skip_special_tokens=True\n",
|
||||
" ).strip()\n",
|
||||
" \n",
|
||||
" base_predictions.append(base_pred)\n",
|
||||
" base_pred = tokenizer.decode(\n",
|
||||
" base_generated_ids[0][base_inputs[\"input_ids\"].shape[1]:],\n",
|
||||
" skip_special_tokens=True\n",
|
||||
" ).strip()\n",
|
||||
" \n",
|
||||
" base_predictions.append(base_pred)\n",
|
||||
" \n",
|
||||
" # Clean up temporary GPU tensors for this step\n",
|
||||
" del base_inputs, base_generated_ids\n",
|
||||
" torch.cuda.empty_cache()\n",
|
||||
" gc.collect()\n",
|
||||
"\n",
|
||||
"print(\"\\nBase Model inferences completed successfully.\")"
|
||||
]
|
||||
|
|
@ -634,105 +646,120 @@
|
|||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import torch\n",
|
||||
"import gc\n",
|
||||
"\n",
|
||||
"video2lora_predictions = []\n",
|
||||
"print(\"Running Video2LoRA Internalization & Inference (zero visual tokens) for all 6 examples...\")\n",
|
||||
"for idx, example in enumerate(examples):\n",
|
||||
" print(f\"Processing Example {idx+1}/{len(examples)}: {example['id']}...\")\n",
|
||||
" \n",
|
||||
" # 1. Reset model to ensure clean base features are extracted\n",
|
||||
" model.reset()\n",
|
||||
" \n",
|
||||
" # 2. Extract context activations\n",
|
||||
" internalize_messages = [\n",
|
||||
" [\n",
|
||||
" {\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"content\": [\n",
|
||||
" {\"type\": \"video\", \"path\": example[\"video_path\"]},\n",
|
||||
" {\"type\": \"text\", \"text\": train_args.internalization_prompt}\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
"with torch.inference_mode():\n",
|
||||
" for idx, example in enumerate(examples):\n",
|
||||
" print(f\"Processing Example {idx+1}/{len(examples)}: {example['id']}...\")\n",
|
||||
" \n",
|
||||
" # 1. Reset model to ensure clean base features are extracted\n",
|
||||
" model.reset()\n",
|
||||
" torch.cuda.empty_cache()\n",
|
||||
" gc.collect()\n",
|
||||
" \n",
|
||||
" # 2. Extract context activations\n",
|
||||
" internalize_messages = [\n",
|
||||
" [\n",
|
||||
" {\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"content\": [\n",
|
||||
" {\"type\": \"video\", \"path\": example[\"video_path\"]},\n",
|
||||
" {\"type\": \"text\", \"text\": train_args.internalization_prompt}\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
" ]\n",
|
||||
" ]\n",
|
||||
" ]\n",
|
||||
" \n",
|
||||
" vlm_inputs = prepare_smolvlm_inputs(\n",
|
||||
" processor,\n",
|
||||
" internalize_messages,\n",
|
||||
" device,\n",
|
||||
" video_fps=train_args.video_fps,\n",
|
||||
" max_frames=train_args.max_frames,\n",
|
||||
" video_size_longest_edge=train_args.video_size_longest_edge,\n",
|
||||
" video_load_backend=train_args.video_load_backend\n",
|
||||
" )\n",
|
||||
" \n",
|
||||
" vlm_inputs = prepare_smolvlm_inputs(\n",
|
||||
" processor,\n",
|
||||
" internalize_messages,\n",
|
||||
" device,\n",
|
||||
" video_fps=train_args.video_fps,\n",
|
||||
" max_frames=train_args.max_frames,\n",
|
||||
" video_size_longest_edge=train_args.video_size_longest_edge,\n",
|
||||
" video_load_backend=train_args.video_load_backend\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" ctx_features, ctx_attn_mask, ctx_position_ids = extract_l2l_fused_text_features(\n",
|
||||
" raw_model,\n",
|
||||
" vlm_inputs,\n",
|
||||
" num_target_layers=model.hypernet.n_layers\n",
|
||||
" )\n",
|
||||
" ctx_features, ctx_attn_mask, ctx_position_ids = extract_l2l_fused_text_features(\n",
|
||||
" raw_model,\n",
|
||||
" vlm_inputs,\n",
|
||||
" num_target_layers=model.hypernet.n_layers\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" # 3. Generate dynamic LoRA weights\n",
|
||||
" generated_loras, _ = model.generate_weights(\n",
|
||||
" ctx_ids=None,\n",
|
||||
" ctx_features=ctx_features,\n",
|
||||
" ctx_attn_mask=ctx_attn_mask,\n",
|
||||
" ctx_position_ids=ctx_position_ids\n",
|
||||
" )\n",
|
||||
" # 3. Generate dynamic LoRA weights\n",
|
||||
" generated_loras, _ = model.generate_weights(\n",
|
||||
" ctx_ids=None,\n",
|
||||
" ctx_features=ctx_features,\n",
|
||||
" ctx_attn_mask=ctx_attn_mask,\n",
|
||||
" ctx_position_ids=ctx_position_ids\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" generated_loras = combine_lora(\n",
|
||||
" generated_loras,\n",
|
||||
" torch.ones(1, dtype=torch.int32, device=device),\n",
|
||||
" lora_bias=model.hypernet.get_head_bias() if model.hypernet.config.use_bias else None\n",
|
||||
" )\n",
|
||||
" \n",
|
||||
" # 4. Patch base layers and inject generated weights\n",
|
||||
" model.patch_lora_forward()\n",
|
||||
" apply_lora_to_layers(\n",
|
||||
" model.base_model,\n",
|
||||
" model.hypernet.layer_indices,\n",
|
||||
" generated_loras,\n",
|
||||
" torch.ones(1, dtype=torch.int32, device=device),\n",
|
||||
" position_ids=None\n",
|
||||
" )\n",
|
||||
" \n",
|
||||
" # 5. Run text-only Q&A inference\n",
|
||||
" prompt_ids = tokenizer.apply_chat_template(\n",
|
||||
" [\n",
|
||||
" {\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"content\": [{\"type\": \"text\", \"text\": example[\"prompt\"]}]\n",
|
||||
" }\n",
|
||||
" ],\n",
|
||||
" tokenize=True,\n",
|
||||
" add_generation_prompt=True,\n",
|
||||
" return_tensors=\"pt\"\n",
|
||||
" ).to(device)\n",
|
||||
" generated_loras = combine_lora(\n",
|
||||
" generated_loras,\n",
|
||||
" torch.ones(1, dtype=torch.int32, device=device),\n",
|
||||
" lora_bias=model.hypernet.get_head_bias() if model.hypernet.config.use_bias else None\n",
|
||||
" )\n",
|
||||
" \n",
|
||||
" # 4. Patch base layers and inject generated weights\n",
|
||||
" model.patch_lora_forward()\n",
|
||||
" apply_lora_to_layers(\n",
|
||||
" model.base_model,\n",
|
||||
" model.hypernet.layer_indices,\n",
|
||||
" generated_loras,\n",
|
||||
" torch.ones(1, dtype=torch.int32, device=device),\n",
|
||||
" position_ids=None\n",
|
||||
" )\n",
|
||||
" \n",
|
||||
" # 5. Run text-only Q&A inference\n",
|
||||
" prompt_ids = tokenizer.apply_chat_template(\n",
|
||||
" [\n",
|
||||
" {\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"content\": [{\"type\": \"text\", \"text\": example[\"prompt\"]}]\n",
|
||||
" }\n",
|
||||
" ],\n",
|
||||
" tokenize=True,\n",
|
||||
" add_generation_prompt=True,\n",
|
||||
" return_tensors=\"pt\"\n",
|
||||
" ).to(device)\n",
|
||||
"\n",
|
||||
" if hasattr(prompt_ids, \"keys\"):\n",
|
||||
" input_ids = prompt_ids[\"input_ids\"]\n",
|
||||
" attention_mask = prompt_ids.get(\"attention_mask\", torch.ones_like(input_ids))\n",
|
||||
" else:\n",
|
||||
" input_ids = prompt_ids\n",
|
||||
" attention_mask = torch.ones_like(input_ids)\n",
|
||||
" if hasattr(prompt_ids, \"keys\"):\n",
|
||||
" input_ids = prompt_ids[\"input_ids\"]\n",
|
||||
" attention_mask = prompt_ids.get(\"attention_mask\", torch.ones_like(input_ids))\n",
|
||||
" else:\n",
|
||||
" input_ids = prompt_ids\n",
|
||||
" attention_mask = torch.ones_like(input_ids)\n",
|
||||
"\n",
|
||||
" generated_ids = model.base_model.generate(\n",
|
||||
" input_ids=input_ids,\n",
|
||||
" attention_mask=attention_mask,\n",
|
||||
" max_new_tokens=train_args.generation_max_new_tokens,\n",
|
||||
" do_sample=False,\n",
|
||||
" pad_token_id=tokenizer.pad_token_id,\n",
|
||||
" eos_token_id=tokenizer.eos_token_id\n",
|
||||
" )\n",
|
||||
" generated_ids = model.base_model.generate(\n",
|
||||
" input_ids=input_ids,\n",
|
||||
" attention_mask=attention_mask,\n",
|
||||
" max_new_tokens=train_args.generation_max_new_tokens,\n",
|
||||
" do_sample=False,\n",
|
||||
" pad_token_id=tokenizer.pad_token_id,\n",
|
||||
" eos_token_id=tokenizer.eos_token_id\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" video2lora_pred = tokenizer.decode(\n",
|
||||
" generated_ids[0][input_ids.shape[1]:],\n",
|
||||
" skip_special_tokens=True\n",
|
||||
" ).strip()\n",
|
||||
" \n",
|
||||
" video2lora_predictions.append(video2lora_pred)\n",
|
||||
" video2lora_pred = tokenizer.decode(\n",
|
||||
" generated_ids[0][input_ids.shape[1]:],\n",
|
||||
" skip_special_tokens=True\n",
|
||||
" ).strip()\n",
|
||||
" \n",
|
||||
" video2lora_predictions.append(video2lora_pred)\n",
|
||||
" \n",
|
||||
" # Clean up temporary GPU tensors for this step\n",
|
||||
" del vlm_inputs, ctx_features, ctx_attn_mask, ctx_position_ids, generated_loras, prompt_ids, input_ids, attention_mask, generated_ids\n",
|
||||
" model.reset()\n",
|
||||
" torch.cuda.empty_cache()\n",
|
||||
" gc.collect()\n",
|
||||
"\n",
|
||||
"# Reset model hooks back to clean base state\n",
|
||||
"model.reset()\n",
|
||||
"torch.cuda.empty_cache()\n",
|
||||
"gc.collect()\n",
|
||||
"\n",
|
||||
"print(\"\\nVideo2LoRA internalization and inferences completed successfully.\")"
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue