Enhance TransformersLLMClient initialization with improved memory management and diagnostics; update Colab notebook execution counts and repository URL

app/council/llm.py

@@ -180,19 +180,50 @@ class TransformersLLMClient:
         cls,
         model_id: str = "google/gemma-4-E4B-it",
         load_in_4bit: bool = True,
-        device_map: Any = "auto",
+        device_map: Any = None,
+        max_memory: dict | None = None,
     ) -> "TransformersLLMClient":
+        """Load Gemma into memory.
+
+        Defaults that work on a single T4 (16 GB):
+          - 4-bit nf4 with double-quant
+          - fp16 compute (T4 has no bf16 silicon)
+          - device_map={"": 0} — force everything to GPU 0; transformers'
+            "auto" sometimes spills the vision tower to CPU on tight cards
+            and then refuses to load with `llm_int8_enable_fp32_cpu_offload`
+            unless given an explicit dict device_map.
+
+        Pass device_map="auto" + max_memory={0: "14GiB", "cpu": "30GiB"}
+        if you need offload (e.g. running on a 12GB card).
+        """
         import torch
         from transformers import AutoModelForCausalLM, AutoProcessor
 
-        # T4 (Turing) doesn't have bf16 hardware — fall back to fp16 there.
+        # Diagnostics so the user sees what they're working with
+        if torch.cuda.is_available():
+            free, total = torch.cuda.mem_get_info()
+            print(
+                f"[load] GPU={torch.cuda.get_device_name(0)} "
+                f"free={free / 1024**3:.2f} GiB / total={total / 1024**3:.2f} GiB "
+                f"bf16={torch.cuda.is_bf16_supported()}"
+            )
+        else:
+            print("[load] no CUDA device — model will load on CPU (very slow)")
+
+        # T4 (Turing) doesn't have bf16 hardware — fall back to fp16.
         compute_dtype = (
             "bfloat16"
             if (torch.cuda.is_available() and torch.cuda.is_bf16_supported())
             else "float16"
         )
 
+        if device_map is None:
+            device_map = {"": 0} if torch.cuda.is_available() else "cpu"
+
         kwargs: dict = {"dtype": compute_dtype, "device_map": device_map}
+        if max_memory is not None:
+            kwargs["max_memory"] = max_memory
+
         if load_in_4bit:
             from transformers import BitsAndBytesConfig
             kwargs["quantization_config"] = BitsAndBytesConfig(
@@ -200,13 +231,19 @@ class TransformersLLMClient:
                 bnb_4bit_compute_dtype=compute_dtype,
                 bnb_4bit_quant_type="nf4",
                 bnb_4bit_use_double_quant=True,
-                # Allow non-quantizable modules (vision tower, embeddings) to
-                # land on CPU when GPU is tight — required on T4 for E4B.
+                # Allows non-quantizable modules (vision tower, embeddings)
+                # to live on CPU if device_map decides to put them there.
                 llm_int8_enable_fp32_cpu_offload=True,
             )
         processor = AutoProcessor.from_pretrained(model_id)
         model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
         model.eval()
+        if torch.cuda.is_available():
+            free, total = torch.cuda.mem_get_info()
+            print(
+                f"[load] done. GPU free={free / 1024**3:.2f} GiB / "
+                f"total={total / 1024**3:.2f} GiB"
+            )
         return cls(model=model, processor=processor)
 
     def complete(

training/train_grpo.ipynb

@@ -25,7 +25,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 1,
    "id": "install",
    "metadata": {},
    "outputs": [],
@@ -43,7 +43,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "id": "clone",
    "metadata": {},
    "outputs": [
@@ -68,7 +68,7 @@
     "\n",
     "if WORKDIR is None:\n",
     "    # 2) Not present — clone. EDIT THIS URL before running on a fresh Colab.\n",
-    "    REPO_URL = \"https://github.com/<your-username>/stocker.git\"\n",
+    "    REPO_URL = \"https://github.com/CRIMSONHydra/stocker.git\"\n",
     "    WORKDIR  = \"/content/stocker\"\n",
     "    assert \"<your-username>\" not in REPO_URL, (\n",
     "        \"Edit REPO_URL in this cell to your fork before running on a fresh runtime.\"\n",
@@ -83,7 +83,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "auth",
    "metadata": {},
    "outputs": [],
@@ -112,14 +112,7 @@
       "  peers: 374 rows\n",
       "[4/6] Loading curated news/forums/macro ...\n",
       "  news: 25 headlines, forums: 17 posts, macro: 12 events\n",
-      "[5/6] Rendering candlestick charts ...\n",
-      "  charts: 126 PNGs in data/charts\n",
-      "[6/6] Done.\n",
-      "task_easy      ticker=AAPL   steps= 43 chart_ok=yes\n",
-      "task_medium    ticker=INTC   steps= 41 chart_ok=yes\n",
-      "task_hard      ticker=META   steps= 42 chart_ok=yes\n",
-      "\n",
-      "All 3 tasks OK.\n"
+      "[5/6] Rendering candlestick charts ...\n"
      ]
     }
    ],
@@ -132,7 +125,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "id": "f55a1108",
    "metadata": {},
    "outputs": [
@@ -140,130 +133,23 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:103: UserWarning: \n",
-      "Error while fetching `HF_TOKEN` secret value from your vault: 'Requesting secret HF_TOKEN timed out. Secrets can only be fetched when running from the Colab UI.'.\n",
-      "You are not authenticated with the Hugging Face Hub in this notebook.\n",
-      "If the error persists, please let us know by opening an issue on GitHub (https://github.com/huggingface/huggingface_hub/issues/new).\n",
-      "  warnings.warn(\n"
+      "[transformers] Current model requires 6192 bytes of buffer for offloaded layers, which seems does not fit any GPU's remaining memory. If you are experiencing a OOM later, please consider using offload_buffers=True.\n"
      ]
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "09427ffe6e7b4be094cef43d13c33fe6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "processor_config.json: 0.00B [00:00, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8874cb7e313545a2b79bad9317858169",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "chat_template.jinja: 0.00B [00:00, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "37bb28cb29c64f97b2699f780174d9ee",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "config.json: 0.00B [00:00, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "459c2365fa5e4cf1ab6897f7a5e9816f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "tokenizer_config.json: 0.00B [00:00, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0f6eab584b334df89d7aa973094c3e58",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "tokenizer.json:   0%|          | 0.00/32.2M [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7c1bbf3af57b4782b0ea600e652b680a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "model.safetensors:   0%|          | 0.00/16.0G [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f5cfa2f28b1844cd9808cf55cfdfc9db",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Loading weights:   0%|          | 0/2076 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "925f4bb541bf49a5b42067ba6e4c10a9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "generation_config.json:   0%|          | 0.00/208 [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Model loaded on: cuda:0\n"
+     "ename": "ValueError",
+     "evalue": "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules in 32-bit, you need to set `llm_int8_enable_fp32_cpu_offload=True` and pass a custom `device_map` to `from_pretrained`. Check https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu for more details. ",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
+      "\u001b[0;32m/tmp/ipykernel_1481/641969956.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;31m# 4-bit BnB by default; drop load_in_4bit=False on L4/A100.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m client = TransformersLLMClient.from_pretrained(\n\u001b[0m\u001b[1;32m      5\u001b[0m     \u001b[0;34m\"google/gemma-4-E4B-it\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m     \u001b[0mload_in_4bit\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m/content/stocker/app/council/llm.py\u001b[0m in \u001b[0;36mfrom_pretrained\u001b[0;34m(cls, model_id, load_in_4bit, device_map)\u001b[0m\n\u001b[1;32m    194\u001b[0m             )\n\u001b[1;32m    195\u001b[0m         \u001b[0mprocessor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mAutoProcessor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_pretrained\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_id\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 196\u001b[0;31m         \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mAutoModelForCausalLM\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_pretrained\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_id\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    197\u001b[0m         \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    198\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcls\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprocessor\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mprocessor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m/usr/local/lib/python3.12/dist-packages/transformers/models/auto/auto_factory.py\u001b[0m in \u001b[0;36mfrom_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, *model_args, **kwargs)\u001b[0m\n\u001b[1;32m    392\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparent_config\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"quantization_config\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    393\u001b[0m                     \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mquantization_config\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mparent_config\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mquantization_config\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 394\u001b[0;31m             return model_class.from_pretrained(\n\u001b[0m\u001b[1;32m    395\u001b[0m                 \u001b[0mpretrained_model_name_or_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mmodel_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mhub_kwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    396\u001b[0m             )\n",
+      "\u001b[0;32m/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py\u001b[0m in \u001b[0;36mfrom_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, use_safetensors, weights_only, fusion_config, disable_mmap, *model_args, **kwargs)\u001b[0m\n\u001b[1;32m   4188\u001b[0m         \u001b[0;31m# Prepare the full device map\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4189\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mdevice_map\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 4190\u001b[0;31m             \u001b[0mdevice_map\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_get_device_map\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_memory\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhf_quantizer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   4191\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4192\u001b[0m         \u001b[0;31m# Finalize model weight initialization\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m/usr/local/lib/python3.12/dist-packages/transformers/integrations/accelerate.py\u001b[0m in \u001b[0;36m_get_device_map\u001b[0;34m(model, device_map, max_memory, hf_quantizer)\u001b[0m\n\u001b[1;32m    371\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    372\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mhf_quantizer\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 373\u001b[0;31m             \u001b[0mhf_quantizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalidate_environment\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice_map\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice_map\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    374\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    375\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mdevice_map\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m/usr/local/lib/python3.12/dist-packages/transformers/quantizers/quantizer_bnb_4bit.py\u001b[0m in \u001b[0;36mvalidate_environment\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     72\u001b[0m             \u001b[0mvalues\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice_map\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     73\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mvalues\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m\"cpu\"\u001b[0m\u001b[0;34m}\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m\"cpu\"\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mvalues\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;34m\"disk\"\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mvalues\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 74\u001b[0;31m                 raise ValueError(\n\u001b[0m\u001b[1;32m     75\u001b[0m                     \u001b[0;34m\"Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     76\u001b[0m                     \u001b[0;34m\"quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;31mValueError\u001b[0m: Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules in 32-bit, you need to set `llm_int8_enable_fp32_cpu_offload=True` and pass a custom `device_map` to `from_pretrained`. Check https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu for more details. "
      ]
     }
    ],