notebooks/performance_and_memory.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration 1 (baseline)\n",
    "\n",
    "* vit-G/14 config from https://arxiv.org/pdf/2103.00020.pdf\n",
    "* BART-L encoder config from https://huggingface.co/facebook/bart-large/blob/main/config.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: yozh did not convert CLIP initialization properly, transformer init is taken from GPT-2\n",
      "Total params: 2.573B\n",
      "ViT params: 1.845B\n",
      "Text transformer params: 0.680B\n",
      "Memory usage (model + grads): 9.6GiB\n"
     ]
    }
   ],
   "source": [
    "from tqdm import trange\n",
    "import torch\n",
    "from clip import CLIP, LeanTransformerConfig\n",
    "\n",
    "\n",
    "vision_config = LeanTransformerConfig(\n",
    "    hidden_size=1664,\n",
    "    num_hidden_layers=48,\n",
    "    num_attention_heads=16,\n",
    "    intermediate_size=8192,\n",
    "    position_embedding_type='absolute',\n",
    "    max_position_embeddings= (336 // 14) ** 2,\n",
    ")\n",
    "\n",
    "text_config = LeanTransformerConfig(\n",
    "    hidden_size=1536,\n",
    "    num_hidden_layers=24,\n",
    "    num_attention_heads=24,\n",
    "    intermediate_size=6144,\n",
    "    position_embedding_type='rotary',\n",
    ")\n",
    "\n",
    "clip = CLIP(embed_dim=1024, image_resolution=336, vision_patch_size=14, context_length=256, vocab_size=30_000,\n",
    "            vision_config=vision_config, text_config=text_config)\n",
    "# note: i could not find output dimension for vit-G/14, extrapolating to 1024 from CLIP paper\n",
    "\n",
    "GRAD_CHECKPOINTS = True\n",
    "# enable gradient checkpointing aka rematerialization\n",
    "if GRAD_CHECKPOINTS:\n",
    "    clip.transformer._get_sequential().gradient_checkpointing = True\n",
    "    clip.visual.transformer._get_sequential().gradient_checkpointing = True\n",
    "\n",
    "\n",
    "GPU_PARAMS_MIXED = True\n",
    "# if enabled, this will emulate a config where (1) gpu params are mostly fp16, but (2) we store fp32 params in RAM\n",
    "if GPU_PARAMS_MIXED:\n",
    "    for param in clip.parameters():\n",
    "        if param.numel() > 2 ** 16:\n",
    "            param.data = param.data.half()\n",
    "\n",
    "\n",
    "for param in clip.parameters():\n",
    "    if param.requires_grad:  # pre-populate grads to avoid fragmentation\n",
    "        param.grad = torch.zeros_like(param)\n",
    "\n",
    "clip = clip.cuda()\n",
    "\n",
    "opt = torch.optim.SGD(clip.parameters(), lr=1e-3)\n",
    "# using SGD as a mock-up for offloading. hivemind will offload optimizer to RAM, so the memory usage will be same as SGD\n",
    "\n",
    "\n",
    "print(f\"Total params: {sum(p.numel() for p in clip.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"ViT params: {sum(p.numel() for p in clip.visual.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Text transformer params: {sum(p.numel() for p in clip.transformer.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Memory usage (model + grads): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Memory usage (no checkpoints, fp32 params) - NO LONGER FITS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Memory usage (with checkpoints, mixed params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:06<00:00,  1.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=1): 10.2GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:17<00:00,  1.75s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=4): 11.7GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:33<00:00,  3.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=8): 13.8GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [01:05<00:00,  6.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=16): 17.9GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for batch_size in (1, 4, 8, 16):\n",
    "    for i in trange(10):\n",
    "        with torch.cuda.amp.autocast():\n",
    "            image = torch.randn(batch_size, 3, 336, 336, device='cuda')\n",
    "            text = torch.randint(30_000, size=(batch_size, 256), device='cuda')\n",
    "            image_features, text_features, tau = clip.forward(image, text)\n",
    "            not_a_real_loss = torch.mean(image_features @ text_features.t() * tau)\n",
    "            not_a_real_loss.backward()\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "            del image, text, image_features, text_features, not_a_real_loss\n",
    "    torch.cuda.synchronize()\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    print(f\"Memory usage (batch={batch_size}): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration 2: PixelFly\n",
    "\n",
    "* use PixelFly block-sparse weight matrices from https://arxiv.org/abs/2112.00029\n",
    "* increase hidden size 1644 -> 2048 to make it compatible with intermediate_size\n",
    "* otherwise same as configuration 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: yozh did not convert CLIP initialization properly, transformer init is taken from GPT-2\n",
      "Total params: 0.820B\n",
      "ViT params: 0.559B\n",
      "Text transformer params: 0.213B\n",
      "Memory usage (model + grads): 6.2GiB\n"
     ]
    }
   ],
   "source": [
    "from tqdm import trange\n",
    "import torch\n",
    "from clip import CLIP, LeanTransformerConfig\n",
    "\n",
    "\n",
    "vision_config = LeanTransformerConfig(\n",
    "    hidden_size=2048,       # <-- changed this line!!!\n",
    "    block_size=64,          # <-- added this line!!!\n",
    "    lowrank_dim=64,         # <-- added this line!!!\n",
    "    num_hidden_layers=48,\n",
    "    num_attention_heads=16,\n",
    "    intermediate_size=8192,\n",
    "    position_embedding_type='absolute',\n",
    "    max_position_embeddings= (336 // 14) ** 2,\n",
    ")\n",
    "\n",
    "text_config = LeanTransformerConfig(\n",
    "    hidden_size=1536,\n",
    "    num_hidden_layers=24,\n",
    "    block_size=96,          # <-- added this line!!!\n",
    "    num_attention_heads=24,\n",
    "    intermediate_size=6144,\n",
    "    position_embedding_type='rotary',\n",
    ")\n",
    "\n",
    "clip = CLIP(embed_dim=1024, image_resolution=336, vision_patch_size=14, context_length=256, vocab_size=30_000,\n",
    "            vision_config=vision_config, text_config=text_config)\n",
    "# note: i could not find output dimension for vit-G/14, extrapolating to 1024 from CLIP paper\n",
    "\n",
    "GRAD_CHECKPOINTS = False\n",
    "# enable gradient checkpointing aka rematerialization\n",
    "if GRAD_CHECKPOINTS:\n",
    "    clip.transformer._get_sequential().gradient_checkpointing = True\n",
    "    clip.visual.transformer._get_sequential().gradient_checkpointing = True\n",
    "\n",
    "\n",
    "GPU_PARAMS_MIXED = False\n",
    "# if enabled, this will emulate a config where (1) gpu params are mostly fp16, but (2) we store fp32 params in RAM\n",
    "if GPU_PARAMS_MIXED:\n",
    "    for param in clip.parameters():\n",
    "        if param.numel() > 2 ** 16:\n",
    "            param.data = param.data.half()\n",
    "\n",
    "\n",
    "for param in clip.parameters():\n",
    "    if param.requires_grad:  # pre-populate grads to avoid fragmentation\n",
    "        param.grad = torch.zeros_like(param)\n",
    "\n",
    "clip = clip.cuda()\n",
    "\n",
    "opt = torch.optim.SGD(clip.parameters(), lr=1e-3)\n",
    "# using SGD as a mock-up for offloading. hivemind will offload optimizer to RAM, so the memory usage will be same as SGD\n",
    "\n",
    "\n",
    "print(f\"Total params: {sum(p.numel() for p in clip.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"ViT params: {sum(p.numel() for p in clip.visual.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Text transformer params: {sum(p.numel() for p in clip.transformer.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Memory usage (model + grads): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Memory usage (no checkpoints, fp32 params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:07<00:00,  1.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=1): 8.2GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:22<00:00,  2.20s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=4): 13.7GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:41<00:00,  4.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=8): 21.0GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for batch_size in (1, 4, 8):\n",
    "    for i in trange(10):\n",
    "        with torch.cuda.amp.autocast():\n",
    "            image = torch.randn(batch_size, 3, 336, 336, device='cuda')\n",
    "            text = torch.randint(30_000, size=(batch_size, 256), device='cuda')\n",
    "            image_features, text_features, tau = clip.forward(image, text)\n",
    "            not_a_real_loss = torch.mean(image_features @ text_features.t() * tau)\n",
    "            not_a_real_loss.backward()\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "            del image, text, image_features, text_features, not_a_real_loss\n",
    "\n",
    "    torch.cuda.synchronize()\n",
    "    torch.cuda.empty_cache()\n",
    "    print(f\"Memory usage (batch={batch_size}): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n",
    "    \n",
    "# note: throughput is a bit slower since we're using a bigger model (1664 -> 2048)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration 3: Sharing with adapters\n",
    "\n",
    "* Share large matrices as in ALBERT: https://arxiv.org/abs/1909.11942\n",
    "* Each layer has non-shared layer-norm and biases as in https://arxiv.org/abs/2107.11817\n",
    "* Each layer has non-shared low-dimensional adapter a-la [LoRA](https://arxiv.org/abs/2106.09685)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: yozh did not convert CLIP initialization properly, transformer init is taken from GPT-2\n",
      "Total params: 0.184B\n",
      "ViT params: 0.089B\n",
      "Text transformer params: 0.048B\n",
      "Memory usage (model + grads): 1.4GiB\n"
     ]
    }
   ],
   "source": [
    "from tqdm import trange\n",
    "import torch\n",
    "from clip import CLIP, LeanTransformerConfig\n",
    "\n",
    "\n",
    "vision_config = LeanTransformerConfig(\n",
    "    hidden_size=1664,\n",
    "    share_large_matrices=True,  # <-- added this line!!!\n",
    "    adapter_dim=32,             # <-- added this line!!!\n",
    "    num_hidden_layers=48,\n",
    "    num_attention_heads=16,\n",
    "    intermediate_size=8192,\n",
    "    position_embedding_type='absolute',\n",
    "    max_position_embeddings= (336 // 14) ** 2,\n",
    ")  # vit-G/14 config from https://arxiv.org/pdf/2103.00020.pdf\n",
    "\n",
    "text_config = LeanTransformerConfig(\n",
    "    hidden_size=1536,\n",
    "    num_hidden_layers=24,\n",
    "    num_attention_heads=24,\n",
    "    intermediate_size=6144,\n",
    "    share_large_matrices=True,  # <-- added this line!!!\n",
    "    adapter_dim=32,             # <-- added this line!!!\n",
    "    position_embedding_type='rotary',\n",
    ")  # BART-L encoder config from https://huggingface.co/facebook/bart-large/blob/main/config.json\n",
    "\n",
    "clip = CLIP(embed_dim=1024, image_resolution=336, vision_patch_size=14, context_length=256, vocab_size=30_000,\n",
    "            vision_config=vision_config, text_config=text_config)\n",
    "# note: i could not find output dimension for vit-G/14, extrapolating to 1024 from CLIP paper\n",
    "\n",
    "GRAD_CHECKPOINTS = False\n",
    "# enable gradient checkpointing aka rematerialization\n",
    "if GRAD_CHECKPOINTS:\n",
    "    clip.transformer._get_sequential().gradient_checkpointing = True\n",
    "    clip.visual.transformer._get_sequential().gradient_checkpointing = True\n",
    "\n",
    "\n",
    "GPU_PARAMS_MIXED = False\n",
    "# if enabled, this will emulate a config where (1) gpu params are mostly fp16, but (2) we store fp32 params in RAM\n",
    "if GPU_PARAMS_MIXED:\n",
    "    for param in clip.parameters():\n",
    "        if param.numel() > 2 ** 16:\n",
    "            param.data = param.data.half()\n",
    "\n",
    "\n",
    "for param in clip.parameters():\n",
    "    if param.requires_grad:  # pre-populate grads to avoid fragmentation\n",
    "        param.grad = torch.zeros_like(param)\n",
    "\n",
    "clip = clip.cuda()\n",
    "\n",
    "opt = torch.optim.SGD(clip.parameters(), lr=1e-3)\n",
    "# using SGD as a mock-up for offloading. hivemind will offload optimizer to RAM, so the memory usage will be same as SGD\n",
    "\n",
    "\n",
    "print(f\"Total params: {sum(p.numel() for p in clip.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"ViT params: {sum(p.numel() for p in clip.visual.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Text transformer params: {sum(p.numel() for p in clip.transformer.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Memory usage (model + grads): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Memory usage (no checkpoints, fp32 params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:06<00:00,  1.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=1): 3.3GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:15<00:00,  1.59s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=4): 8.1GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:29<00:00,  2.95s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=8): 14.5GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:42<00:00,  4.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=12): 20.8GiB\n"
     ]
    }
   ],
   "source": [
    "for batch_size in (1, 4, 8, 12):\n",
    "    for i in trange(10):\n",
    "        with torch.cuda.amp.autocast():\n",
    "            image = torch.randn(batch_size, 3, 336, 336, device='cuda')\n",
    "            text = torch.randint(30_000, size=(batch_size, 256), device='cuda')\n",
    "            image_features, text_features, tau = clip.forward(image, text)\n",
    "            not_a_real_loss = torch.mean(image_features @ text_features.t() * tau)\n",
    "            not_a_real_loss.backward()\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "            del image, text, image_features, text_features, not_a_real_loss\n",
    "\n",
    "    torch.cuda.synchronize()\n",
    "    torch.cuda.empty_cache()\n",
    "    print(f\"Memory usage (batch={batch_size}): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n",
    "    \n",
    "# note: throughput is a bit slower since we're using a bigger model (1664 -> 2048)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration 4: reversible\n",
    "\n",
    "* same as configuration 3, but now using reversible layers instead of checkpoints\n",
    "* same reversible strategy as in [Reformer](https://arxiv.org/abs/2001.04451) (using [revlib](https://github.com/clashluke/revlib) under the hood)\n",
    "* reversible transformer keeps 2 sets of activations in memory instead of keeping one checkpoint for every layer\n",
    "* as a result, we can get a __batch size of up to 64__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: yozh did not convert CLIP initialization properly, transformer init is taken from GPT-2\n",
      "Total params: 0.184B\n",
      "ViT params: 0.089B\n",
      "Text transformer params: 0.048B\n",
      "Memory usage (model + grads): 1.4GiB\n"
     ]
    }
   ],
   "source": [
    "from tqdm import trange\n",
    "import torch\n",
    "from clip import CLIP, LeanTransformerConfig\n",
    "\n",
    "\n",
    "vision_config = LeanTransformerConfig(\n",
    "    reversible=True,            # <-- added this line!!!\n",
    "    hidden_size=1664,\n",
    "    share_large_matrices=True,  # <-- added this line!!!\n",
    "    adapter_dim=32,             # <-- added this line!!!\n",
    "    num_hidden_layers=48,\n",
    "    num_attention_heads=16,\n",
    "    intermediate_size=8192,\n",
    "    position_embedding_type='absolute',\n",
    "    max_position_embeddings= (336 // 14) ** 2,\n",
    ")  # vit-G/14 config from https://arxiv.org/pdf/2103.00020.pdf\n",
    "\n",
    "text_config = LeanTransformerConfig(\n",
    "    hidden_size=1536,\n",
    "    num_hidden_layers=24,\n",
    "    num_attention_heads=24,\n",
    "    intermediate_size=6144,\n",
    "    reversible=True,            # <-- added this line!!!\n",
    "    share_large_matrices=True,  # <-- added this line!!!\n",
    "    adapter_dim=32,             # <-- added this line!!!\n",
    "    position_embedding_type='rotary',\n",
    ")  # BART-L encoder config from https://huggingface.co/facebook/bart-large/blob/main/config.json\n",
    "\n",
    "clip = CLIP(embed_dim=1024, image_resolution=336, vision_patch_size=14, context_length=256, vocab_size=30_000,\n",
    "            vision_config=vision_config, text_config=text_config)\n",
    "# note: i could not find output dimension for vit-G/14, extrapolating to 1024 from CLIP paper\n",
    "\n",
    "# note: gradient checkpoints are not used if model is reversible!\n",
    "\n",
    "for param in clip.parameters():\n",
    "    if param.requires_grad:  # pre-populate grads to avoid fragmentation\n",
    "        param.grad = torch.zeros_like(param)\n",
    "\n",
    "clip = clip.cuda()\n",
    "\n",
    "opt = torch.optim.SGD(clip.parameters(), lr=1e-3)\n",
    "# using SGD as a mock-up for offloading. hivemind will offload optimizer to RAM, so the memory usage will be same as SGD\n",
    "\n",
    "\n",
    "print(f\"Total params: {sum(p.numel() for p in clip.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"ViT params: {sum(p.numel() for p in clip.visual.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Text transformer params: {sum(p.numel() for p in clip.transformer.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Memory usage (model + grads): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:07<00:00,  1.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=1): 1.8GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:19<00:00,  1.91s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=4): 2.4GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [01:08<00:00,  6.89s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=16): 4.8GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [04:32<00:00, 27.24s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=64): 14.2GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for batch_size in (1, 4, 16, 64):\n",
    "    for i in trange(10):\n",
    "        with torch.cuda.amp.autocast():\n",
    "            image = torch.randn(batch_size, 3, 336, 336, device='cuda')\n",
    "            text = torch.randint(30_000, size=(batch_size, 256), device='cuda')\n",
    "            image_features, text_features, tau = clip.forward(image, text)\n",
    "            not_a_real_loss = torch.mean(image_features @ text_features.t() * tau)\n",
    "            not_a_real_loss.backward()\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "            \n",
    "            del image, text, image_features, text_features, not_a_real_loss\n",
    "\n",
    "    torch.cuda.synchronize()\n",
    "    torch.cuda.empty_cache()\n",
    "    print(f\"Memory usage (batch={batch_size}): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n",
    "\n",
    "# note the singificantly larger batch size (up to 64 per gpu and still some vram left)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration 5: going crazy\n",
    "\n",
    "\n",
    "* combine configurations 1-4\n",
    "* increase model size as much as possible\n",
    "* this setup only makes sense if you have 50+ collaborators with a 2080Ti or better GPU each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: yozh did not convert CLIP initialization properly, transformer init is taken from GPT-2\n",
      "Total params: 0.605B\n",
      "ViT params: 0.183B\n",
      "Text transformer params: 0.140B\n",
      "Memory usage (model + grads): 4.5GiB\n"
     ]
    }
   ],
   "source": [
    "from tqdm import trange\n",
    "import torch\n",
    "from clip import CLIP, LeanTransformerConfig\n",
    "\n",
    "\n",
    "vision_config = LeanTransformerConfig(\n",
    "    hidden_size=8192,            # <-- changed this line!!!\n",
    "    share_large_matrices=True,   # <-- added this line!!!\n",
    "    block_size=64,               # <-- added this line!!!\n",
    "    lowrank_dim=120,             # <-- added this line!!!\n",
    "    adapter_dim=8,               # <-- added this line!!!\n",
    "    num_hidden_layers=64,        # <-- changed this line!!!\n",
    "    num_attention_heads=64,      # <-- changed this line!!!\n",
    "    intermediate_size=32768,     # <-- changed this line!!!\n",
    "    reversible=True,             # <-- added this line!!!\n",
    "    position_embedding_type='absolute',\n",
    "    max_position_embeddings= (336 // 14) ** 2,\n",
    ")   # let's call this vit-ludicrous /14\n",
    "\n",
    "text_config = LeanTransformerConfig(\n",
    "    hidden_size=8192,            # <-- changed this line!!!\n",
    "    share_large_matrices=True,   # <-- added this line!!!\n",
    "    block_size=64,               # <-- added this line!!!\n",
    "    lowrank_dim=120,             # <-- added this line!!!\n",
    "    adapter_dim=8,               # <-- added this line!!!\n",
    "    num_hidden_layers=64,        # <-- changed this line!!!\n",
    "    num_attention_heads=64,      # <-- changed this line!!!\n",
    "    intermediate_size=32768,     # <-- changed this line!!!\n",
    "    reversible=True,             # <-- added this line!!!\n",
    "    position_embedding_type='rotary',\n",
    ")\n",
    "\n",
    "clip = CLIP(embed_dim=4096, image_resolution=336, vision_patch_size=14, context_length=256, vocab_size=30_000,\n",
    "            vision_config=vision_config, text_config=text_config)\n",
    "# note: embed dim is now 4096, mostly for lulz\n",
    "\n",
    "\n",
    "GPU_PARAMS_MIXED = False\n",
    "# if enabled, this will emulate a config where (1) gpu params are mostly fp16, but (2) we store fp32 params in RAM\n",
    "if GPU_PARAMS_MIXED:\n",
    "    for param in clip.parameters():\n",
    "        if param.numel() > 2 ** 16:\n",
    "            param.data = param.data.half()\n",
    "\n",
    "\n",
    "for param in clip.parameters():\n",
    "    if param.requires_grad:  # pre-populate grads to avoid fragmentation\n",
    "        param.grad = torch.zeros_like(param)\n",
    "\n",
    "clip = clip.cuda()\n",
    "\n",
    "opt = torch.optim.SGD(clip.parameters(), lr=1e-3)\n",
    "# using SGD as a mock-up for offloading. hivemind will offload optimizer to RAM, so the memory usage will be same as SGD\n",
    "\n",
    "\n",
    "print(f\"Total params: {sum(p.numel() for p in clip.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"ViT params: {sum(p.numel() for p in clip.visual.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Text transformer params: {sum(p.numel() for p in clip.transformer.parameters()) / 1e9 :.3f}B\")\n",
    "print(f\"Memory usage (model + grads): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:56<00:00,  5.65s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=1): 5.7GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [01:34<00:00,  9.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=2): 6.5GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [03:03<00:00, 18.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=4): 8.3GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [05:46<00:00, 34.69s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage (batch=8): 11.9GiB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for batch_size in (1, 2, 4, 8):\n",
    "    for i in trange(10):\n",
    "        with torch.cuda.amp.autocast():\n",
    "            image = torch.randn(batch_size, 3, 336, 336, device='cuda')\n",
    "            text = torch.randint(30_000, size=(batch_size, 256), device='cuda')\n",
    "            image_features, text_features, tau = clip.forward(image, text)\n",
    "            not_a_real_loss = torch.mean(image_features @ text_features.t() * tau)\n",
    "            not_a_real_loss.backward()\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "            \n",
    "            del image, text, image_features, text_features, not_a_real_loss\n",
    "\n",
    "    torch.cuda.synchronize()\n",
    "    torch.cuda.empty_cache()\n",
    "    print(f\"Memory usage (batch={batch_size}): {torch.cuda.max_memory_allocated() / 2 ** 30:.1f}GiB\")\n",
    "\n",
    "# note: technically speaking, batch 16 fits into memory with mixed params, but it takes a minute for each run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py38",
   "language": "python",
   "name": "py38"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}