g

scripts/cxrvlm_colab_train.ipynb

@@ -409,6 +409,24 @@
    ],
    "id": "cell-resume"
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "cell-smoke-md"
+   },
+   "source": "## 6.0 Pre-train smoke test (loads model, exercises both stages)\n\nQuick sanity check before kicking off real training. Verifies:\n\n1. **Library versions** — torch / transformers / peft / bitsandbytes / accelerate compatible (a mismatched bitsandbytes is the #1 cause of QLoRA failures).\n2. **Model loads** — Vicuna-7B 4-bit + RAD-DINO + projection + LoRA construct without error.\n3. **Both stages compile** — `set_stage1_mode()` and `set_stage2_mode()` each produce a finite loss on a dummy batch.\n4. **VRAM fits** — prints peak VRAM after each stage's forward.\n\nRuns in ~2 min on L4/3090. If anything fails here, the real train cell will fail too — fix it now rather than 15 min into stage 1. After this cell, GPU is freed so the train subprocess starts with a clean slate.",
+   "id": "cell-smoke-md"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "cell-smoke"
+   },
+   "source": "# ─── 0. Library versions + GPU capability ──────────────────────────────\nimport torch, sys, importlib, gc\nprint(f'Python              : {sys.version.split()[0]}')\nprint(f'torch               : {torch.__version__}  (cuda={torch.version.cuda})')\nfor pkg in ('transformers', 'peft', 'bitsandbytes', 'accelerate',\n            'omegaconf', 'datasets', 'tokenizers'):\n    try:\n        m = importlib.import_module(pkg)\n        print(f'{pkg:<20}: {getattr(m, \"__version__\", \"?\")}')\n    except ImportError as e:\n        print(f'{pkg:<20}: NOT INSTALLED ({e})')\n\nassert torch.cuda.is_available(), 'CUDA not available — training will be unusable.'\ngpu = torch.cuda.get_device_properties(0)\nbf16_ok = torch.cuda.is_bf16_supported()\nprint(f'\\nGPU                 : {gpu.name}  ({gpu.total_memory/1e9:.1f} GB)')\nprint(f'Compute capability  : {gpu.major}.{gpu.minor}  (BF16 supported: {bf16_ok})')\n\n# Sanity-check the bf16/fp16 config matches what the GPU supports.\nif train_cfg.training.bf16 and not bf16_ok:\n    print('\\n⚠️  train_cfg.training.bf16=True but this GPU is pre-Ampere — '\n          'switch to fp16 profile in cell-cfg (e.g. T4 profile).')\nif train_cfg.training.fp16 and train_cfg.training.bf16:\n    print('\\n⚠️  Both fp16 and bf16 are True. Trainer will pick bf16; fp16 ignored.')\n\n# ─── 1. Build model ────────────────────────────────────────────────────\nfrom model import CXRVisionLanguageModel\nprint('\\nBuilding model (Vicuna-7B 4-bit + RAD-DINO + LoRA)...')\n_t0 = __import__('time').time()\nsmoke_model = CXRVisionLanguageModel(model_cfg).cuda()\nprint(f'Model built in {__import__(\"time\").time()-_t0:.1f}s')\n\ndef _vram(tag):\n    used = torch.cuda.memory_allocated()/1e9\n    peak = torch.cuda.max_memory_allocated()/1e9\n    print(f'  VRAM [{tag:<22}]  used={used:5.2f} GB   peak={peak:5.2f} GB')\n\ndef _count_trainable(m):\n    tp = sum(p.numel() for p in m.parameters() if p.requires_grad)\n    ap = sum(p.numel() for p in m.parameters())\n    return tp, ap\n\n_vram('after model load')\n\n# ─── 2. Build a dummy batch ────────────────────────────────────────────\n# Tiny input matching the model's forward signature. Image placeholder\n# token id 32000 will be replaced by 32 visual tokens internally.\nIMG_TOKEN_ID = 32000\nB = 1\ndummy_images   = torch.randn(B, 3, model_cfg.image_encoder.img_size,\n                                   model_cfg.image_encoder.img_size,\n                             device='cuda', dtype=torch.float32)\n# 'USER: <image> hello ASSISTANT: hi' style — short prompt is fine for smoke test.\ndummy_ids = torch.tensor([[1, 11889, 29901, IMG_TOKEN_ID, 22172, 319, 1799,\n                           9047, 13566, 29901, 7251]],\n                         device='cuda', dtype=torch.long)\ndummy_mask   = torch.ones_like(dummy_ids)\ndummy_labels = dummy_ids.clone()\n\n# ─── 3. Stage 1 smoke ──────────────────────────────────────────────────\ntorch.cuda.reset_peak_memory_stats()\nsmoke_model.set_stage1_mode()\ntp, ap = _count_trainable(smoke_model)\nprint(f'\\n[Stage 1] trainable params: {tp:,} / {ap:,}  ({100*tp/ap:.3f}%)')\nwith torch.autocast('cuda',\n                    dtype=torch.bfloat16 if train_cfg.training.bf16 else torch.float16):\n    out1 = smoke_model(images=dummy_images, input_ids=dummy_ids,\n                       attention_mask=dummy_mask, labels=dummy_labels)\nloss1 = out1['loss'].item()\nprint(f'[Stage 1] dummy loss = {loss1:.4f}  (finite: {torch.isfinite(out1[\"loss\"]).item()})')\nassert torch.isfinite(out1['loss']).item(), 'Stage 1 loss is NaN/Inf — investigate before training.'\nout1['loss'].backward()    # ensure backward also works under the chosen precision\nprint('[Stage 1] backward() OK')\n_vram('stage 1 fwd+bwd')\n\n# ─── 4. Stage 2 smoke ──────────────────────────────────────────────────\nsmoke_model.zero_grad(set_to_none=True)\ntorch.cuda.reset_peak_memory_stats()\nsmoke_model.set_stage2_mode()\ntp, ap = _count_trainable(smoke_model)\nprint(f'\\n[Stage 2] trainable params: {tp:,} / {ap:,}  ({100*tp/ap:.3f}%)')\nwith torch.autocast('cuda',\n                    dtype=torch.bfloat16 if train_cfg.training.bf16 else torch.float16):\n    out2 = smoke_model(images=dummy_images, input_ids=dummy_ids,\n                       attention_mask=dummy_mask, labels=dummy_labels)\nloss2 = out2['loss'].item()\nprint(f'[Stage 2] dummy loss = {loss2:.4f}  (finite: {torch.isfinite(out2[\"loss\"]).item()})')\nassert torch.isfinite(out2['loss']).item(), 'Stage 2 loss is NaN/Inf — investigate before training.'\nout2['loss'].backward()\nprint('[Stage 2] backward() OK')\n_vram('stage 2 fwd+bwd')\n\n# ─── 5. Cleanup — free GPU for the real train subprocess ───────────────\ndel smoke_model, out1, out2, dummy_images, dummy_ids, dummy_mask, dummy_labels\ngc.collect(); torch.cuda.empty_cache()\nprint(f'\\n✅ Smoke test passed. VRAM after cleanup: '\n      f'{torch.cuda.memory_allocated()/1e9:.2f} GB used.')",
+   "execution_count": null,
+   "outputs": [],
+   "id": "cell-smoke"
+  },
   {
    "cell_type": "markdown",
    "metadata": {