LLM_Techniques_and_utils/Traverse_through_sub-modules_of_PyTorch_Model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🧠 Useful Util method when working with Language Models - \n",
    "\n",
    "### For traversing through nested attributes (or sub-modules) of a PyTorch module\n",
    "\n",
    "Language models generally have multiple layers/modules that are nested within each other and you might want to replace a specific layer of a pre-trained model with a new one, or you might want to freeze the weights of specific layers while fine-tuning the rest, or you might want to inspect the parameters or output of a specific layer for debugging. \n",
    "\n",
    "For all those cases you need a way to pinpoint the exact sub-module you want to work with.\n",
    "\n",
    "### And the `get_child_module_by_names()` util method does just that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🧠 useful Utility method when working with Language Models\n",
    "# Traverse through nested attributes (or sub-modules) of a PyTorch module\n",
    "\n",
    "import torch\n",
    "from typing import Iterable\n",
    "from torch.nn import Module\n",
    "\n",
    "def get_child_module_by_names(module: Module, names: Iterable[str]) -> Module:\n",
    "    obj = module\n",
    "    for getter in map(lambda name: lambda obj: getattr(obj, name), names):\n",
    "        obj = getter(obj)\n",
    "    # outer lambda function takes `name` as an arg and returns another lambda function.\n",
    "    # inner lambda function takes `obj` as its arg and applies `getattr(obj, name)`.\n",
    "    return obj\n",
    "\n",
    "######### USING THE ABOVE UTIL METHOD #################\n",
    "# Let's look at an example where we might want to freeze\n",
    "# the embedding layers of a pre-trained BERT model for fine-tuning.\n",
    "# This would involve setting the `requires_grad` attribute of the parameters\n",
    "# of these layers to `False`.\n",
    "\n",
    "from transformers import BertModel\n",
    "\n",
    "def get_child_module_by_names(module, names):\n",
    "    obj = module\n",
    "    for getter in map(lambda name: lambda obj: getattr(obj, name), names):\n",
    "        obj = getter(obj)\n",
    "    return obj\n",
    "\n",
    "# Load a pre-trained BERT model\n",
    "model = BertModel.from_pretrained('bert-base-uncased')\n",
    "\n",
    "# The names to the embedding layer in BERT\n",
    "names_to_embeddings = ['embeddings']\n",
    "\n",
    "# Use the function to fetch the embeddings module\n",
    "embeddings = get_child_module_by_names(model, names_to_embeddings)\n",
    "\n",
    "# Freeze the embeddings module\n",
    "for param in embeddings.parameters():\n",
    "    param.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lets see how the `get_child_module_by_names()` method  is working\n",
    "\n",
    "👉 The first argument `module`, is a PyTorch `nn.Module` object. It could be any object but, in the context of HuggingFace language models, it's likely to be a model like BERT, GPT-NeoX, etc. The second argument, `names`, is an iterable (most likely a list or tuple) containing the names of the nested sub-modules we're trying to access, in order.\n",
    "\n",
    "👉 The function initializes `obj` as `module`, which means initially `obj` is the outermost module. \n",
    "\n",
    "----------------\n",
    "\n",
    "👉 Next, the function enters a loop over `names`. For each name in the `names` list, it applies a function that takes an object and returns the attribute of that object corresponding to the current name. The lambda function `lambda name: lambda obj: getattr(obj, name)` is a higher-order function that returns another function. This nested function will, when called with an object, try to fetch the attribute with the name `name`.\n",
    "\n",
    "This style is an example of nested lambda function within the map function.\n",
    "\n",
    "📌 The outer lambda function takes `name` as an argument and returns another lambda function.\n",
    "\n",
    "📌 The inner lambda function takes `obj` as its argument and applies `getattr(obj, name)`.\n",
    "\n",
    "`getattr()` takes two args: the object you want to get an attribute from, and a string that is the name of the attribute you want to get. For example, if we have an object `foo` with an attribute `bar`, we can get `bar` by calling `getattr(foo, 'bar')`.\n",
    "\n",
    "This can be useful when we need to access attributes dynamically, i.e., we don't know which attributes we need to access until the program is running. `getattr` is used here to dynamically access the sub-modules of a PyTorch module, given the string names of these sub-modules.\n",
    "\n",
    "👉 The loop can be read like this: \"For each `name` in `names`, get me the attribute of `obj` that has the name `name`, and set `obj` to be this newly fetched attribute\". This way, with each iteration, `obj` becomes a deeper sub-module within the original module.\n",
    "\n",
    "👉 The function repeats this until all names in `names` have been processed. The final object, `obj`, returned by this function will be the innermost requested module (the attribute of the module represented by the last name in `names`).\n",
    "\n",
    "\n",
    "---------------\n",
    "\n",
    "👉 Use of the method\n",
    "\n",
    "The code in the image has an example use-case, where we might want to freeze the embedding layers of a pre-trained BERT model for fine-tuning. This would involve setting the `requires_grad` attribute of the parameters of these layers to `False`. \n",
    "\n",
    "The specific path to the embeddings module (`['embeddings']`) is determined by the structure of the BERT model and can be found in its documentation or source code.\n",
    "\n",
    "The parameters of the fetched embeddings module are iterated over, and their `requires_grad` attribute is set to `False`, which effectively \"freezes\" them. This means that when we backpropagate through the network during fine-tuning, the gradients w.r.t these parameters will not be computed, and the optimizer will not update these parameters.\n",
    "\n",
    "--------\n",
    "\n",
    "👉 In the `get_child_module_by_names()` method why we needed the nested lambda structure.\n",
    "\n",
    "📌 The nested lambda structure in `get_child_module_by_names` is a way to dynamically create a sequence of getter functions, each tailored to access a specific attribute of an object. Its designed to navigate through nested sub-modules of a PyTorch model by sequentially accessing named attributes.\n",
    "\n",
    "--------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why use nested lambda structure for getting sub-modules of a Pytorch Language Model\n",
    "\n",
    "📌  In the `get_child_module_by_names()` method why we needed the nested lambda structure.\n",
    "\n",
    "📌 The nested lambda structure in `get_child_module_by_names` is a way to dynamically create a sequence of getter functions, each tailored to access a specific attribute of an object. Its designed to navigate through nested sub-modules of a PyTorch model by sequentially accessing named attributes.\n",
    "\n",
    "-------\n",
    "\n",
    "📌 Example code in image:\n",
    "\n",
    "Consider a PyTorch model with a nested structure, like `model.layer1.sublayer2.attribute`. If you want to access `attribute`, you must first navigate through `layer1` and then `sublayer2`.\n",
    "\n",
    "📌 Here's how the nested lambda works in this context:\n",
    "\n",
    "1. **Outer Lambda**: Takes the name of a sub-module (e.g., 'layer1') and returns an inner lambda function.\n",
    "\n",
    "2. **Inner Lambda**: This function is configured to take an object (`obj`) and apply `getattr(obj, 'layer1')`, effectively fetching `obj.layer1`.\n",
    "\n",
    "📌 Applied to the sequence of names:\n",
    "\n",
    "- For 'layer1', you get a function that does `getattr(obj, 'layer1')`.\n",
    "- For 'sublayer2', you get another function for `getattr(obj, 'sublayer2')`.\n",
    "- These functions are then applied in sequence to dive deeper into the model's structure.\n",
    "\n",
    "📌 In this simple example, the nested lambda functions act as a chain of getters, each fetching the next level in the nested structure until the final attribute is reached. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nested lambda structure for getting\n",
    "# sub-modules of a Pytorch Language Model\n",
    "\n",
    "# Imagine a nested object structure\n",
    "\n",
    "class SubLayer:\n",
    "    def __init__(self):\n",
    "        self.attribute = \"value\"\n",
    "\n",
    "class Layer:\n",
    "    def __init__(self):\n",
    "        self.sublayer = SubLayer()\n",
    "\n",
    "model = Layer()\n",
    "\n",
    "# Names to access the nested attribute\n",
    "names = [\"sublayer\", \"attribute\"]\n",
    "\n",
    "# The method in action\n",
    "obj = model\n",
    "for name in names:\n",
    "    # Equivalent to: obj = getattr(obj, name)\n",
    "    getter = lambda obj, name=name: getattr(obj, name)\n",
    "    obj = getter(obj)\n",
    "\n",
    "print(obj)  # Outputs: value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------\n",
    "\n",
    "💡BONUS DISCUSSION - How the `getter` in the `get_child_module_by_names()` function works ❓\n",
    "\n",
    "The `getter` in the `get_child_module_by_names()` function is an example of a higher-order function, which is a function that returns another function. In this case, `getter` is generated by the expression `lambda name: lambda obj: getattr(obj, name)`. This can be broken down into two parts:\n",
    "\n",
    "1. `lambda name: ...` is a function that takes a `name` and returns another function.\n",
    "\n",
    "2. `... lambda obj: getattr(obj, name)` is the function that gets returned. It's a function that takes an `obj` and uses `getattr` to fetch the attribute of `obj` that has the name `name`.\n",
    "\n",
    "So, the `getter` is a function that, when given an object, fetches the attribute of that object with a certain name. The specific name is determined when the getter is created.\n",
    "\n",
    "Here's a simple example of how the  `getter` works:\n",
    "\n",
    "In this example, `create_getter()` is analogous to `lambda name: lambda obj: getattr(obj, name)` in the `get_child_module_by_names()` function, and `getter1` and `getter2` are analogous to `getter`.\n",
    "\n",
    "This example demonstrates the power of higher-order functions. By creating different getters, we can fetch different attributes from the same object, or the same attribute from different objects, depending on how we use them.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "```py\n",
    "# Define a simple class with a couple of attributes\n",
    "class MyClass:\n",
    "    attr1 = 'hello'\n",
    "    attr2 = 'world'\n",
    "\n",
    "# Create an instance of the class\n",
    "my_instance = MyClass()\n",
    "\n",
    "# Define a function that creates a getter for an attribute\n",
    "def create_getter(name):\n",
    "    return lambda obj: getattr(obj, name)\n",
    "\n",
    "# Create a getter for 'attr1'\n",
    "getter1 = create_getter('attr1')\n",
    "\n",
    "# Use the getter to fetch 'attr1' from my_instance\n",
    "print(getter1(my_instance))  # Output: hello\n",
    "\n",
    "# Create a getter for 'attr2'\n",
    "getter2 = create_getter('attr2')\n",
    "\n",
    "# Use the getter to fetch 'attr2' from my_instance\n",
    "print(getter2(my_instance))  # Output: world\n",
    "\n",
    "\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}