Motivation

In case of compute bound, instead of using a large model, can we just fine-tune a small model for domain-specific examples and still get good results? The goal is to fine-tune an SLM for tool calling and check the results. Although current SLMs are already better at tool calling, I still wanted to learn more about SFT and tool calling hands-on.

Model & Dataset

• The latest SLMs from Qwen are already performing very well on tool calling, so I chose the earlier base model for my experiment, which is Qwen2.5-0.5B.
• I have used the xLAM dataset from Salesforce, where a sample looks like:

{
  "query": "Find the sum of all the multiples of 3 and 5 between 1 and 1000. Also find the product of the first five prime numbers.",
  "tools": [
    {
      "name": "math_toolkit.sum_of_multiples",
      "description": "Find the sum of all multiples of specified numbers within a specified range.",
      "parameters": {
        "lower_limit": {
          "type": "int",
          "description": "The start of the range (inclusive).",
          "required": True
        },
        "upper_limit": {
          "type": "int",
          "description": "The end of the range (inclusive).",
          "required": True
        },
        "multiples": {
          "type": "list",
          "description": "The numbers to find multiples of.",
          "required": True
        }
      }
    },
    {
      "name": "math_toolkit.product_of_primes",
      "description": "Find the product of the first n prime numbers.",
      "parameters": {
        "count": {
          "type": "int",
          "description": "The number of prime numbers to multiply together.",
          "required": True
        }
      }
    }
  ],
  "answers": [
    {
      "name": "math_toolkit.sum_of_multiples",
      "arguments": {
        "lower_limit": 1,
        "upper_limit": 1000,
        "multiples": [3, 5]
      }
    },
    {
      "name": "math_toolkit.product_of_primes",
      "arguments": {
        "count": 5
      }
    }
  ]
}

• I have used True instead of true, since the sample is shown as a Python dict here rather than raw JSON.
• Each JSON object specifies:
  • query — the query with the params possibly needed for the tools
  • tools — the available tools which can be called, where each tool has a description and specifies its arguments; each argument has a description, the type of arg it takes, and whether it’s required
  • answers — a list of tools and their corresponding args (key-value pairs) which need to be passed to the tool

Preparing the Dataset

• Instead of using all 60k samples, I took only 2,500 samples, where I used 2000 as train, 100 as validation, and 400 as test samples.
• Given the query, tools, and answers, I needed to apply a chat template to convert them into the required format.
• For that purpose, I used the chat template of Qwen2.5-0.5B and defined two functions — one for train/validation, and another for test.

def generate_train_val_sample(sample):
  messages = [
      {
          "role": "user",
          "content": sample["query"]
      },
      {
          "role": "assistant",
          "tool_calls": [
              {
                  "function": {
                      "name": answer["name"],
                      "arguments": answer["arguments"]
                  }
              }
              for answer in json.loads(sample["answers"])
          ]
      }
  ]

  tools = json.loads(sample["tools"])
  text = tokenizer.apply_chat_template(
        messages,
        tools=tools,
        tokenize=False,
        add_generation_prompt=False,
    )

  return  {"text": text}

def generate_test_sample(sample):
  messages = [
      {
          "role": "user",
          "content": sample["query"]
      }
  ]

  tools = json.loads(sample["tools"])
  text = tokenizer.apply_chat_template(
        messages,
        tools=tools,
        tokenize=False,
        add_generation_prompt=True,
    )

  return  {"text": text}

• As can be seen from the functions above, we pass answers in the case of train and validation and set add_generation_prompt to False, since we are not generating — we are comparing against the already-known text.

• For test, we don’t pass answers, and we set add_generation_prompt to True.

• Using these two functions, we can convert any sample into the required format.

Converted train/validation sample

<|im_start|>system
You are a helpful assistant.

# Tools

You may call one or more functions to assist with the user query.

You are provided with function signatures within <tools></tools> XML tags:
<tools>
{"name": "get_order", "description": "Fetches the order information for a given order ID using the Toolbench RapidAPI.", "parameters": {"is_id": {"description": "The ID of the order to be fetched.", "type": "str", "default": ""}}}
</tools>

For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
<tool_call>
{"name": <function-name>, "arguments": <args-json-object>}
</tool_call><|im_end|>
<|im_start|>user
Retrieve the order status for \'ORD11223\' and \'ORD44556\'.<|im_end|>
<|im_start|>assistant
<tool_call>
{"name": "get_order", "arguments": {"is_id": "ORD11223"}}
</tool_call>
<tool_call>
{"name": "get_order", "arguments": {"is_id": "ORD44556"}}
</tool_call><|im_end|>'

Converted test sample

<|im_start|>system
You are a helpful assistant.

# Tools

You may call one or more functions to assist with the user query.

You are provided with function signatures within <tools></tools> XML tags:
<tools>
{"name": "get_order", "description": "Fetches the order information for a given order ID using the Toolbench RapidAPI.", "parameters": {"is_id": {"description": "The ID of the order to be fetched.", "type": "str", "default": ""}}}
</tools>

For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
<tool_call>
{"name": <function-name>, "arguments": <args-json-object>}
</tool_call><|im_end|>
<|im_start|>user
Retrieve the order status for \'ORD11223\' and \'ORD44556\'.<|im_end|>
<|im_start|>assistant

• I’ve used the same example above to show how it looks in both train/validation and test form.
• As we can see, the text contains all the required pieces — system prompt, user query, and available tools — and depending on the type of sample, it either includes the assistant’s response and ends with the <|im_end|> token, or it ends right at <|im_start|>assistant, triggering the model to generate the response.

SFT Setup

• I used PEFT and TRL for doing supervised fine-tuning, and didn’t use Unsloth because Unsloth requires some of the latest GPUs, and my current laptop (RTX 1050, 3GB) doesn’t have that.
• I had the option of using Colab, but at the same time, I wanted to check the internal working of the functions I was importing, so instead I switched to my local setup.

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    task_type=TaskType.CAUSAL_LM,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    modules_to_save=["lm_head", "embed_tokens"],
    ensure_weight_tying=True,
)

training_args = SFTConfig(
    # Training schedule / optimization
    per_device_train_batch_size = BATCH_SIZE,      # Batch size per GPU
    gradient_accumulation_steps = 4,      # Effective batch size = 1 * 4 = 4
    warmup_steps = WARMUP_STEPS,
    max_steps=MAX_STEPS,
    learning_rate = 2e-4,                 # Learning rate for the optimizer
    optim = "paged_adamw_8bit",           # Optimizer

    # Logging / reporting
    eval_strategy="steps",
    eval_steps=EVAL_STEPS,
    save_strategy="steps",
    save_steps=SAVE_STEPS,
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    logging_steps=1,                      # Log training metrics every N steps
    report_to="trackio",                  # Experiment tracking tool
    run_name="testing",
    output_dir=OUTPUT_DIR,                # Where to save model checkpoints and logs

    max_length=1024,                      # Maximum input sequence length
    activation_offloading=True,           # Offload activations to CPU to reduce GPU memory usage
)

Training Logs

Evaluation

• first i checked whether the predicted answers have all the available tools as present in ground truth and then checked for the arguments present and the values passed against those arguments
• compared the result of finetuned base model with instruct version of qwen-2.5-0.5B

Results

Before fine-tuning (base model)

User Query: I need an advanced chart for 'BINANCE:ETHUSDT' with MACD and RSI indicators, in light theme, and a 4h interval. Also, provide a list of stocks that are overbought according to RSI on page 3.
Actual Output: [{"name": "advanced_chart", "arguments": {"symbol": "BINANCE:ETHUSDT", "studies": "MACD,RSI", "theme": "light", "interval": "4h"}}, {"name": "rsi_overbought", "arguments": {"page": 3}}]
Predicted Output: To fetch the advanced chart for 'BINANCE:ETHUSDT' with MACD and RSI indicators, in light theme, and a 4h interval, you can use the following code:

{
    "name": "advanced_chart",
    "arguments": {
        "symbol": "BINANCE:ETHUSDT",
        "studies": ["MACD", "RSI"],
        "format": "png",
        "theme": "light",
        "interval": "4h",
        "width": 800,
        "height": 600,
        "timezone": "Etc/UTC"
    }
}

This code will fetch the advanced chart for 'BINANCE:ETHUSDT' with MACD and RSI indicators, in light theme, and a 4h interval. The arguments include the symbol, studies, format, theme, interval, width, height, and timezone.

To fetch a list of stocks that are overbought according to RSI on page 3, you can use the following code:

{
    "name": "rsi_overbought",
    "arguments": {
        "page": 3,
        "studies": ["RSI"],
        "format": "png",
        "theme": "light",
        "interval": "4h",
        "width": 800,
        "height": 600,
        "timezone": "Etc/UTC"
    }
}

This code will fetch a list of stocks that are overbought according to RSI on page 3, with MACD and RSI indicators, in light theme, and a 4h interval. The arguments include the page number, studies, format, theme, interval, width, height, and timezone.

The base model, before fine-tuning, doesn’t call tools at all — it responds conversationally and hallucinates extra arguments (format, width, height, timezone) that were never part of the tool’s actual schema.

After fine-tuning

User Query: I need an advanced chart for 'BINANCE:ETHUSDT' with MACD and RSI indicators, in light theme, and a 4h interval. Also, provide a list of stocks that are overbought according to RSI on page 3.
Actual Output: [{"name": "advanced_chart", "arguments": {"symbol": "BINANCE:ETHUSDT", "studies": "MACD,RSI", "theme": "light", "interval": "4h"}}, {"name": "rsi_overbought", "arguments": {"page": 3}}]
Predicted Output: <tool_call>
{"name": "advanced_chart", "arguments": {"symbol": "BINANCE:ETHUSDT", "studies": "MACD,RSI", "theme": "light", "interval": "4h"}}
</tool_call>
<tool_call>
{"name": "rsi_overbought", "arguments": {"page": 3}}
</tool_call>

After fine-tuning, the model correctly emits both tool calls in the expected <tool_call> format, with arguments matching the ground truth exactly.

What I Learned

• For SLMs, the embedding and lm_head are generally kept tied, to reduce the parameter count.
• In this case, the tokenizer’s vocab size was lesser than the number of tokens the embedding module could actually represent — this extra headroom becomes useful if new tokens need to be added during fine-tuning.
• When adding extra tokens, the embedding and lm_head also need to be trained — otherwise the model produces weird/garbage tokens as output.

Code

Full training and evaluation code: https://github.com/riteshrm/qwen2.5-tool-calling-lora

Motivation

Perspective of a distribution via “Mass Conservation”

Tip

• 
• Where:
  • is temporal evolution of density at time t at a particular location

Flow Models

Earlier attempts to learn the velocity

Flow Matching

Issues with directly using Flow Matching

Warning

• As we can see in above figure, two vector fields are intersecting each other, which means our vector field approximator will learn somewhat average of these two
• And at inference time we will not be able to go to or , we will go somewhere in between.
• And also there is no linearity as it was in case of score based models, where we can deterministically compute the linear part

Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class FlowMatching(nn.Module):
    def __init__(self, velocity_predictor):
        super().__init__()

        self.velocity_predictor = velocity_predictor

    def flow_lerp(self, x_1):
        bsz = x_1.shape[0]
        t = torch.rand((bsz,)).to(x_1.device)
        t_broadcast = t.view(bsz, 1, 1, 1)
        x_0 = torch.randn_like(x_1)
        # the flow-matching lerp
        x_t = (1 - t_broadcast) * x_0 + t_broadcast * x_1
        return x_0, x_t, t
    
    @torch.no_grad()
    def sample(self, shape, num_steps):
        bsz = shape[0]
        x_t = torch.randn(shape, device=self.velocity_predictor.device)
        delta_t = torch.tensor([1/num_steps] * bsz).to(x_t.device).view([bsz, *([1] * len(x_t.shape[1:]))])
        for i in range(num_steps):
            t = torch.tensor([i/num_steps] * bsz).to(x_t.device)
            predicted_velocity = self.velocity_predictor(x_t, t)
            x_t = x_t + predicted_velocity * delta_t
        return x_t
    
    def loss(self, x_1):
        x_0, x_t, t = self.flow_lerp(x_1)
        target_velocity = x_1-x_0
        predicted_velocity = self.velocity_predictor(x_t, t)
        return torch.nn.functional.mse_loss(target_velocity, predicted_velocity)

Motivation

• Objective: Given a complex data distribution how can we go from low probability to high probability
• Solution: Use gradient, which will guide to go towards high density
• So if we somehow have access to the gradient of i.e. data then we can use it.
• 
• But according to property of probabilities
  • which is intractable
  • Thus is also intractable and also numerically unstable in low density regions
• So instead of that we can consider
  • It is tractable
    • 
    • since is independent of x so it can be treated as constant.
  • It Points in the same direction as
    • 
  • Numerically stable
    • Since gradient depends on ratio of two numbers which are in similar range
    • Similar range because if is small then there is possibility that p is also small
• So we can follow the score i.e gradient of with respect to x to move towards data points
• It is possible that following the scores, we might always end up with data points which has high probability density and thus we might not be able to move to data points of low probability.

• But we don’t have , so we try to estimate the score with

• which gives:

  • 

• But we don’t have access to

• There are various attempts to estimate the score, but let’s focus on more common one Denoising Score Matching

Denoising Score Matching

• The idea is to add noise to data, since we can get score of gaussian
  • 
• Therefore:
  • 
  • 
• Adding noise to distribution
  • 
  • where is perturbation kernel
  • It is basically take all the points from data sample and transform them into noisy samples where tells how much noise you are adding
• Estimating the score of the noised distribution
  • Definition of score matching
    • 
  • From some maths
    • 
    • where is tractable and equal to
• We now have a tractable loss
• However in , is not exactly
  • let’s say we used :
    • will be close to but poor estimation in low density regions.
  • if we use :
    • will be far from but good estimations in low density regions.
• So how do we achieve best of both worlds
  • The idea is to have score matching with varying noise
  • 
  • 
  • which gives loss as:
    • . NCSN: Noise Conditioned Score Network
• Now we got the score approximator as , how can we use this to generate samples

• But in above cases, first of all we need to define T, which will be responsible for discretizing the noising process.
• And also if somehow we can make this process evolve in continuous time instead of discrete ones then we can use differential equations efficiently.

From Discrete to Continuous

• Lets make discrete DDPM into continuous one
• 
• We want a differential evolution in x as function f everything else.
• 
• Lets make as function of t
  • 
  • where
    • : Noise we add between two discrete steps
    • : rate of noise at t
    • : Increment of time
• Now replace with , which gives
  • 
• If -> 0, then (Differential Evolution)
• Using Taylor expansion and wiener process, we can approximate first and second term respectively
  • 
  • 
• 
• 
  • Where
    • is drift coefficient which is deterministic
    • is diffusion coefficient which is stochastic
• SDE Variants

NoteCommon SDE Variants

Variance Preserving	Variance Exploding
Ex: DDPM	Ex:NCSN

• Generalised training objective
  • 
    • where
• Training a generalised score model
  1. Sample , and
  2. Create a noisy sample
     • . where and can be defined based on the variant of SDE
  3. Use and to predict via
  4. Compute Loss and backpropagate
     • 
• Inference of trained model
  • Forward process: Perturb data into noise
    • 
  • Reverse process: noise back to data
    • . This derived using Fokker-Planck equation
  • Intuition of above equation
    • : forward drift coefficient
    • : forward diffusion coefficient
    • : Wiener process but different than forward one
    • : correction to drift coefficient
    • Basically in forward process we diffused from data to noise and to go back we need some corrections.
• Sampling using Euler-Maruyama
  1. Sample Noise i.e. 
  2. Use Euler-Maruyama to go through reverse SDE
     • 
  3. Obtain final sample
• getting
  • 
  • So
• Limitations of SDE
  • Stochastic term as means:
    • Slower solver: Need 1000-2000 steps
    • More source of error: Discretization as well as injected stochastic noise

Probability Flow ODE

• Hypothetically write SDE as an ODE
  • , where is null
  • No stochastic term means:
    • Faster solver: leverage mathematical properties to go fast vs slow
    • Focused source of error: Now we will have error from discretization only
• Lets derive ODE
• Forward SDE:
• Probability flow of x at t using Fokker-Planck Equation on forward SDE gives:
  • 
• After applying some algebraic operations, we get continuity equation as:
  • 
  • where
    • is velocity of probability flow
• So we got PF-ODE as and
  • where:
    • and are modeling assumptions
    • is score approximated by the model
• In above formulation, the Probability Flow is preserved between reverse SDE and reverse ODE but the trajectories are different

Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class VPSDE(nn.Module):
    def __init__(self, beta_min, beta_max, eps, score_model):
        super().__init__()
        self.beta_min = beta_min
        self.beta_max = beta_max
        self.eps = eps
        self.score_model = score_model
    
    def get_integrated_beta_t(self, t):
        return self.beta_min*t + 0.5 * (self.beta_max - self.beta_min) * t**2
    
    def get_beta_t(self, t):
        return self.beta_min + (self.beta_max - self.beta_min) * t

    def get_integrated_alpha_t(self, t):
        return torch.exp(-0.5 * self.get_integrated_beta_t(t))
    
    def get_integrated_sigma_t(self, t):
        return torch.sqrt(1 - torch.exp(-self.get_integrated_beta_t(t)))
    
    def perturbation_kernel(self, x_0, t, noise=None):
        alpha_t = self.get_integrated_alpha_t(t).view(-1, 1, 1, 1)
        sigma_t = self.get_integrated_sigma_t(t).view(-1, 1, 1, 1)
        if noise is None:
            noise = torch.randn_like(x_0)
        x_t = alpha_t * x_0 + sigma_t * noise
        return x_t, noise, sigma_t
    
    def sde(self, x, t):
        drift = -0.5 * self.get_beta_t(t).view(-1, 1, 1, 1) * x
        diffusion = torch.sqrt(self.get_beta_t(t)).view(-1, 1, 1, 1)
        return drift, diffusion
    
    def sample(self, shape, y, num_steps):
        x_t = torch.randn(shape)
        ts = torch.linspace(1, self.eps, num_steps)
        delta_t = -1/num_steps
        for t in ts:
            score = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
            drift, diffusion = self.sde(x_t, t)
            x_prevt = x_t + (drift - diffusion**2*score) * delta_t + diffusion * torch.sqrt(-delta_t)*torch.randn_like(x_t)
            x_t = x_prevt
        return x_t
    
    def sample_pf_ode(self, shape, y, num_steps):
        x_t = torch.randn(shape)
        ts = torch.linspace(1, self.eps, num_steps)
        delta_t = -1/num_steps
        for t in ts:
            score = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
            drift, diffusion = self.sde(x_t, t)
            x_prevt = x_t + (drift - 0.5*diffusion**2*score) * delta_t
            x_t = x_prevt
        return x_t

    def loss(self, batch):
        x_0, y = batch
        t = self.eps + (1- self.eps)*torch.rand((x_0.shape[0], ))
        x_t, noise, sigma_t = self.perturbation_kernel(x_0, t)
        s_theta = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
        loss = F.mse_loss(s_theta, -noise/sigma_t, reduction='mean')
        return loss

import torch
import torch.nn as nn
import torch.nn.functional as F

class subVPSDE(nn.Module):
    def __init__(self, beta_min, beta_max, eps, score_model):
        super().__init__()
        self.beta_min = beta_min
        self.beta_max = beta_max
        self.eps = eps
        self.score_model = score_model
    
    def get_integrated_beta_t(self, t):
        return self.beta_min*t + 0.5 * (self.beta_max - self.beta_min) * t**2
    
    def get_beta_t(self, t):
        return self.beta_min + (self.beta_max - self.beta_min) * t

    def get_integrated_alpha_t(self, t):
        return torch.exp(-0.5 * self.get_integrated_beta_t(t))
    
    def get_integrated_sigma_t(self, t):
        return 1 - torch.exp(-self.get_integrated_beta_t(t))
    
    def perturbation_kernel(self, x_0, t, noise=None):
        alpha_t = self.get_integrated_alpha_t(t).view(-1, 1, 1, 1)
        sigma_t = self.get_integrated_sigma_t(t).view(-1, 1, 1, 1)
        if noise is None:
            noise = torch.randn_like(x_0)
        x_t = alpha_t * x_0 + sigma_t * noise
        return x_t, noise, sigma_t
    
    def sde(self, x, t):
        drift = (-0.5 * self.get_beta_t(t) * x).view(-1, 1, 1, 1)
        diffusion = torch.sqrt(self.get_beta_t(t)*(1-torch.exp(-2*self.get_integrated_beta_t(t)))).view(-1, 1, 1, 1)
        return drift, diffusion
    
    def sample(self, shape, y, num_steps):
        x_t = torch.randn(shape)
        ts = torch.linspace(1, self.eps, num_steps)
        delta_t = -1/num_steps
        for t in ts:
            score = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
            drift, diffusion = self.sde(x_t, t)
            x_prevt = x_t + (drift - diffusion**2*score) * delta_t + diffusion * torch.sqrt(-delta_t)*torch.randn_like(x_t)
            x_t = x_prevt
        return x_t
    
    def sample_pf_ode(self, shape, y, num_steps):
        x_t = torch.randn(shape)
        ts = torch.linspace(1, self.eps, num_steps)
        delta_t = -1/num_steps
        for t in ts:
            score = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
            drift, diffusion = self.sde(x_t, t)
            x_prevt = x_t + (drift - 0.5*diffusion**2*score) * delta_t
            x_t = x_prevt
        return x_t

    def loss(self, batch):
        x_0, y = batch
        t = self.eps + (1- self.eps)*torch.rand((x_0.shape[0], ))
        x_t, noise, sigma_t = self.perturbation_kernel(x_0, t)
        s_theta = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
        loss = F.mse_loss(s_theta, -noise/sigma_t, reduction='mean')
        return loss

import torch
import torch.nn as nn
import torch.nn.functional as F

class VESDE(nn.Module):
    def __init__(self, sigma_min, sigma_max, eps, score_model):
        super().__init__()
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
        self.score_model = score_model
        self.eps = eps
    
    def get_sigma_t(self, t):
        return self.sigma_min*((self.sigma_max/self.sigma_min) **t)
    
    def perturbation_kernel(self, x_0, t, noise=None):
        sigma_t = self.get_sigma_t(t).view(-1, 1, 1, 1)
        if noise is None:
            noise = torch.randn_like(x_0)
        x_t = x_0 + sigma_t * noise
        return x_t, noise, sigma_t
    
    def sde(self, x, t):
        drift =torch.zeros_like(x)
        diffusion = (self.get_sigma_t(t) * torch.sqrt(2*torch.log(self.sigma_max/self.sigma_min))).view(-1, 1, 1, 1)
        return drift, diffusion
    
    def sample(self, shape, y, num_steps):
        x_t = torch.randn(shape)
        ts = torch.linspace(1, self.eps, num_steps)
        delta_t = -1/num_steps
        for t in ts:
            score = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
            drift, diffusion = self.sde(x_t, t)
            x_prevt = x_t + (drift - diffusion**2*score) * delta_t + diffusion * torch.sqrt(-delta_t)*torch.randn_like(x_t)
            x_t = x_prevt
        return x_t
    
    def sample_pf_ode(self, shape, y, num_steps):
        x_t = torch.randn(shape)
        ts = torch.linspace(1, self.eps, num_steps)
        delta_t = -1/num_steps
        for t in ts:
            score = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
            drift, diffusion = self.sde(x_t, t)
            x_prevt = x_t + (drift - 0.5*diffusion**2*score) * delta_t
            x_t = x_prevt
        return x_t

    def loss(self, batch):
        x_0, y = batch
        t = self.eps + (1- self.eps)*torch.rand((x_0.shape[0], ))
        x_t, noise, sigma_t = self.perturbation_kernel(x_0, t)
        s_theta = self.score_model(x_t, t.view(-1, 1, 1, 1), y)
        loss = F.mse_loss(s_theta, -noise/sigma_t, reduction='mean')
        return loss

Denoising Diffusion Probabilistic Models

• Denoising Diffusion Probabilistic Models operate within a variational framework, much like VAEs

• DDPMs introduce a clever twist that tackles some of the challenges faced by their predecessors

• DDPMs involve two most distinct stochastic processes:-

• The forward Pass (Fixed Encoder):

  • Gradually corrupts data by injecting gaussian noise over multiple steps via a transition kernel
  • The data evolves into an isotropic gaussian distribution, effectively becoming pure noise
  • Encoder is fixed not learned.

• The reverse denoising process(Learnable Decoder):
  • A neural net learns to reverse the noise corruption through a parameterized distribution
  • Starting from pure noise, this process iteratively denoises to generate realistic samples
  • Each individual denoising step is a more manageable task than generating a complete sample from scratch as VAEs often attempt to do

Forward Process

• Forward process involves taking the noisy image from previous timestep and inject more noise
• Forward transition kernel: where is output, is mean and is variance
• Given , In paper and
• 
• Noise injection form:
• lets say
• Then
• Lets say

Symbol	Meaning
	noise schedule
	gaussian noise

• 
• 
• if we substitute value of in above equation then
• 
• 
• 

Tip

• When we merge two normal distributions with different variance, the new distribution is as per below formula
  • 

• Using above intuition we can change accordingly the forward equation
• 
• 
• 
• If we generalize above equation it becomes
• 

Reverse Process

Reverse Transition (Parameterization)

• 
• Two neural networks to parameterize normal distribution and
• In paper variance is fixed so we only need to predict mean
• We want to maximize the log likelihood of our data as we do in case of VAEs
• 
• but it is not easily computable as it depends on all the other :-
• So keeping track of other T-1 random variables is not possible as a solution we can compute ELBO as we did in VAEs

Objective: Log-Likelihood and ELBO

• lets say when we have and and > for all x then if we can maximize then will also be maximized

• Using the above analogy, lets say and

• From above equations, we can see and are almost same, just there is one difference the KL-Divergence term which is being subtracted so if we are subtracting something then it is lesser than and thus if we can somehow maximize then we can maximize the likelihood of data samples

• 

• Now we have to maximize

• Now our loss becomes negative of log likelihood

• Even now our loss still contains and to eliminate that we have to make some changes to KL term

From KL Term to VLB (Step-by-step Expansion)

• Now lets focus on KL-term only
  • =
  • the denominator in the above expression can be replaced using bayes’ theorem
  • 
  • Numerator can be replaced by using joint probability.
  • Now
  • where can be written as full joint distribution of entire diffusion trajectory
  • Now putting back the denominator value in KL term as gives us
  • 
  • If we expand above equation as:-
  • +
• lets add above derived kl term in Equation 1 which gives
• 
• Notice gets cancelled with
• Now our loss function becomes:
• Which gives VLB

VLB Expansion (Product to Summation)

• As we can see in VLB numerator is basically forward pass and denominator can be written as:
• and is simply normal distribution with mean 0 and std 1
• We can also write numerator as product of probabilities as
• Now our VLB becomes
• 
• Now take out which gives
• 
• Now we can transform product to summation which gives
• 
• Now from summation one take out t=1 formulation and keep everything as it is which gives
• 
• now the numerator in summation of logs can be written as:
• 

Conditioning Trick: Reduce Variance with

• Now all the three terms 2 in numerator and 1 in denominator have really high variance, since we dont know what we started from that’s why author tried to reduce the variance by conditioning on x_0 and thus it reduces the choice of possible candidates. Intuitively given a noisy sample how can you say that what part do i have to remove, whereas if we also have x_0 then we might guess what part of noise to remove.

• So

• Notice If we have not taken out the t=1 term then there have been a loop as

• 

• lets put back the value of back in original equation which gives

• 

• lets expand summation one

• 

• the 3rd term in above equation if expanded is just

• Now our loss looks like

• 

• lets open up the 3rd and 4th term over log

• $-log(p(x_T)) + {t=2}^{T}log + log(q(x_T x_0)) - log(q(x_1 x_0)) + log(q(x_1 x_0)) - log(p{}(x_0 x_1)) $

• which gives us finally

• 

Final Loss Decomposition (, , )

• Let’s say
  • 
  • 
  • 
• We can neglect since it does not involve any parameters
• For
• Recall when we applied bayes’ rule, due to that both forward and reverse process are in same form now
• Now lets focus on and
• 
• since is fixed which is
• 
• also has a closed form solution which is equal to
• 
• For derivation check out What are Diffusion Models?
• But for now assume
  • 
  • 
  • Since _t is fixed, we can neglect this

Noise Prediction Reparameterization (Predict epsilon)

• from forward pass put the value of X_0 in
• 
• from above equation
• 
• which gives
• 
• Our model needs to predict
• Since is available as input at training time, we can reparameterize the gaussian noise term instead to predict from input at time step
• 
• 
• now comparing and
• 
• Authors from the DDPM paper neglected the weighting term citing that it improves the sample quality and implementation is quite easy.
• So
• Now for
• Authors decided to neglect the last term but basically it was comparing with , like how probable is the predicted given
• By neglecting , while generating sample, we add noise till but at , we just return the sample

Sampling Step ( from )

• Now back to starting sample generation
• , which gives
• 
• now putting the value of gives
• 

Implementation

• Below is simple implementation of forward process, reverse process, sampling and loss, where eps_theta is our denoiser model
• For the complete class-conditioned DDPM implementation, check out:- nanoDiffusion/ddpm.py

import torch
class Diffusion():
    def __init__(self, beta_start, beta_end, timesteps):
        self.beta_start = beta_start
        self.beta_end = beta_end
        self.timesteps = timesteps
        self.beta_t = torch.linspace(self.beta_start, self.beta_end, self.timesteps)
        self.alpha_t = 1- self.beta_t
        self.alphabar_t = torch.cumprod(self.alpha_t, dim=0)
        self.eps_theta = torch.nn.Module()

    def forward_process(self, x_0, eps, t):
        alphabar_t = self.alphabar_t[t].view(-1, 1, 1, 1)
        x_t = torch.sqrt(alphabar_t)*x_0 + torch.sqrt(1 - alphabar_t)*eps
        return x_t
    
    def reverse_process(self, x_t, t, eps=None):
        pred_eps = self.eps_theta(x_t, t)
        alphabar_t = self.alphabar_t[t].view(-1, 1, 1, 1)
        alpha_t = self.alpha_t[t].view(-1, 1, 1, 1)
        beta_t = self.beta_t[t].view(-1, 1, 1, 1)
        if eps is None:
            x_prevt = (1/torch.sqrt(alpha_t))*(x_t - (beta_t/torch.sqrt(1 - alphabar_t))*pred_eps)
        else:
            x_prevt = (1/torch.sqrt(alpha_t))*(x_t - (beta_t/torch.sqrt(1 - alphabar_t))*pred_eps) + torch.sqrt(beta_t)*eps
        return x_prevt
    
    def sample(self, shape):
        x_t = torch.randn(shape)
        
        for t in range(self.timesteps-1, -1, -1):
            if t==0:
                x_t = self.reverse_process(x_t, t)
            else:
                eps = torch.randn(shape)
                x_t = self.reverse_process(x_t, t, eps)
        return x_t
    
    def loss(self, x_0):
        t = torch.randint(0, self.timesteps, (x_0.shape[0],))
        eps = torch.randn_like(x_0)
        x_t = self.forward_process(x_0, eps, t)
        pred_eps = self.eps_theta(x_t, t)
        return torch.nn.functional.mse_loss(eps, pred_eps)

---

Faster Sampling (Denoising Diffusion Implicit Models)

Implementation

• For the complete class-conditioned DDIM implementation, check out:- nanoDiffusion/ddim.py

import torch
class Diffusion():
    def __init__(self, beta_start, beta_end, timesteps):
        self.beta_start = beta_start
        self.beta_end = beta_end
        self.timesteps = timesteps
        self.beta_t = torch.linspace(self.beta_start, self.beta_end, self.timesteps)
        self.alpha_t = 1- self.beta_t
        self.alphabar_t = torch.cumprod(self.alpha_t, dim=0)
        self.eps_theta = torch.nn.Module()

    def forward_process(self, x_0, eps, t):
        alphabar_t = self.alphabar_t[t].view(-1, 1, 1, 1)
        x_t = torch.sqrt(alphabar_t)*x_0 + torch.sqrt(1 - alphabar_t)*eps
        return x_t
    
    def reverse_process(self, x_t, t, eps=None):
        pred_eps = self.eps_theta(x_t, t)
        alphabar_t = self.alphabar_t[t].view(-1, 1, 1, 1)
        alpha_t = self.alpha_t[t].view(-1, 1, 1, 1)
        beta_t = self.beta_t[t].view(-1, 1, 1, 1)
        if eps is None:
            x_prevt = (1/torch.sqrt(alpha_t))*(x_t - (beta_t/torch.sqrt(1 - alphabar_t))*pred_eps)
        else:
            x_prevt = (1/torch.sqrt(alpha_t))*(x_t - (beta_t/torch.sqrt(1 - alphabar_t))*pred_eps) + torch.sqrt(beta_t)*eps
        return x_prevt
    
    def sample_ddim(self, shape, num_steps):
        x_t = torch.randn(shape)

        steps = torch.linspace(self.timesteps-1, 0, num_steps, dtype=torch.int)

        for i in range(num_steps):
            t = steps[i]
            alphabar_t = self.alphabar_t[t].view(-1, 1, 1, 1)
            pred_eps = self.eps_theta(x_t, t)
            x_hat = (x_t - (1-alphabar_t).sqrt()*pred_eps)/alphabar_t.sqrt()
            if i == len(steps)-1:
                return x_hat
            prev_t = steps[i+1]
            alphabar_prevt = self.alphabar_t[prev_t].view(-1, 1, 1, 1)
            x_prevt = alphabar_prevt.sqrt()*x_hat + (1-alphabar_prevt).sqrt()*pred_eps
            x_t = x_prevt
    
    def loss(self, x_0):
        t = torch.randint(0, self.timesteps, (x_0.shape[0],))
        eps = torch.randn_like(x_0)
        x_t = self.forward_process(x_0, eps, t)
        pred_eps = self.eps_theta(x_t, t)
        return torch.nn.functional.mse_loss(eps, pred_eps)

Autoencoders

• They are used for:-
  • converting higher dimension to lower dimension
  • Representation Learning

What We Actually Want?

• Instead of encoding to data points, encode into some distribution so that we can generate new data points.
• But even if we try to encode to a distribution, the model will again try to memorise things as it just have to minimize the reconstruction loss.
• Get variations in reconstructed images.
• So we want is that:
  • There should be some sort of latent distribution which is just not encoded to points, it should be wide enough to capture entire latent space so that we can have higher variance and thus variation in reconstructed images.
  • And also distributions should not be far apart from each other so that if we sample any latent then it should give us something meaningful not random images.

How Do We Achieve Above Requirements?

• Gaussian distribution would help in achieving above objectives
  • It will try to have mean as close to zero as possible thus keep distributions closer to each other.
  • It will try to have unit variance thus ensuring that distributions have higher variance and hence disable the network to collapse distributions to a point.
• So basically what we want is that our encoded distribution should look like a Gaussian distribution.
• We can use some sort of metric that will help in keeping the encoded distribution and Gaussian distribution as similar as possible.
• And that metric is KL-Divergence, which we can use to do so.

Before Moving Forward Some Prerequisites

Let’s Define Our Model

Let’s Do Some Maths

Introducing

• 
• x 1 Multiply by 1
• . Since . because summation of all the z that could have produced x is 1, which is what we are trying to calculate. You can take analogy of discrete example also if we consider only z which could have produced x.
• bring inside integral since it does not depend on z and can be treated as constant.
• 
  • basically above term is equal to but we want to somehow introduce z because is independent of z and using is constant because we can take constant outside and the remaining integral will give 1.
• Apply the equation
• Multiply by 1
• Split the expectation
• By definition of KL Divergence
• where ELBO is evidence lower bound and second term is KL divergence since it will always be greater than 0
• Since we are trying to maximize the likelihood then even if we maximize the first term i.e. ELBO that will work for us since second term will always be greater than 0
• For example
  • Total Compensation = Base Salary + Bonus()
  • Total Compensation Base Salary
• Following the above analogy, we can also say that
  • 
  • So if we maximize the right side then left side will also be maximised

Rewriting ELBO

• following chain rule of probability
• Split the expectation
• Using KL Divergence
• So to maximize the whole term we need to maximize the first one and minimize the second term
• Profit = Revenue - Costs
• To get maximum profit, we need to maximize revenue at the same time minimize the cost

What Maximizing ELBO Means

• Maximizing ELBO means:

  • Maximizing the first term: maximizing the reconstruction likelihood of the decoder
  • Minimizing the second term: minimizing the distance between learned distribution and the prior belief we have over the latent variable

• When we want to maximize a function then we move in the direction of gradient

• When we want to minimize a function then we move against the direction of gradient

• expectation over

Note

In theory we maximize ELBO. In code, we usually minimize the negative ELBO, which becomes reconstruction loss plus KL loss.

---

Need of Reparameterization

Reparameterization Trick

• So we need a new estimator:- Let’s just move the source of randomness outside the model
• Instead of directly sampling from , we rewrite sampling as
  • z = where

• expectation over
• 
• 
• 
• 

Code

• Below code is the simple implementation of how a VAE looks like in code
• Notice in case of encoder, we predict log of variance because neural network might give negative numbers also.

import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 196),
            nn.Tanh(),
            nn.Linear(196, 48),
            nn.Tanh()
        )

        self.mu_fc = nn.Sequential(
            nn.Linear(48, 16),
            nn.Tanh(),
            nn.Linear(16, 2)
        )
        self.logvar_fc = nn.Sequential(
            nn.Linear(48, 16),
            nn.Tanh(),
            nn.Linear(16, 2)
        )
        self.decoder = nn.Sequential(
            nn.Linear(2, 16),
            nn.Tanh(),
            nn.Linear(16, 48),
            nn.Tanh(),
            nn.Linear(48, 196),
            nn.Tanh(),
            nn.Linear(196, 28*28),
            nn.Tanh()
        )
    
    def forward(self, x):
        ##Encoding
        mu, logvar = self.encode(x)

        ##Sampling

        z = self.sample(mu, logvar)

        ## Decoding
        x = self.decode(z)
        return mu, logvar, x
    
    def encode(self, x):
        x = x.reshape(x.shape[0], 28*28)
        x = self.encoder(x)
        mu, logvar = self.mu_fc(x), self.logvar_fc(x)
        return mu, logvar
    
    def decode(self, x):
        x = self.decoder(x)
        x = x.reshape(x.shape[0], 1, 28, 28)
        return x
    
    def sample(self, mu, log_var):
        rand = torch.randn_like(mu)
        std = torch.sqrt(torch.exp(log_var))
        return rand * std + mu

def kl_loss(mu, log_var):
    return torch.mean(0.5 * torch.sum(torch.exp(log_var) + mu**2 - 1 - log_var, dim=-1))

def recon_loss(pred, target):
    return F.mse_loss(pred, target)

Why VAEs Produce Blurry Samples

• Multiple possible images can correspond to similar latent representations. When many plausible images map to nearby latent regions, the decoder tends to generate an averaged reconstruction.
• This overlap occurs partly because of KL regularization, where the encoder is forced to keep latent representations close to a simple prior distribution (usually a Gaussian). As a result, different images may occupy overlapping regions in latent space.
• Additionally, VAEs typically assume the decoder output follows a Gaussian distribution:
  • 
• Due to this assumption, maximizing likelihood becomes equivalent to minimizing Mean Squared Error (MSE) between the real and predicted image.
• Now suppose two possible pixel values for a region are 0 and 1 for similar latent vectors. Under MSE, the optimal prediction becomes their average:
  • 
• Therefore, instead of producing one sharp possibility, the decoder learns to predict averaged pixel values across multiple plausible outputs, leading to blurry images.

Summary

• Encoder learns
• Decoder learns
• KL-divergence keeps latent distribution close to prior distribution
• Reconstruction loss helps preserve input information
• Reparameterization helps backpropagate through sampling

Why Relative Position Embedding?

• It captures how two tokens are positioned with respect to each other

• For Example:-

  •   my daughter called her brother last night  [Sentence:-1]
      0     1 (n)  2 (m)  3    4      5     6

  • Absolute Position Embedding of token daughter depends on position of token i.e. 1 (n)
  • Relative Position Embedding of token ‘daughter’ with respect to token ‘called’ -1 (n-m)
    • Where -1 or in general it will tell whether token ‘daughter’ is before or after ‘called’ and at what distance.

• Query at position m attends Key at position n and to calculate attention score, it depends more on relative position than absolute position.

•   last night my daughter called her brother  [Sentence:-2]
    0      1    2    3        4    5     6

• In above sentence even though the APE is different compared to previous sentence their RPE is still the same

• In case of APE, there is linear relationship between tokens at pos and pos + k, and we expect our model to learn this relationship

  • = *

• In case of RPE, we are explicitly telling our model that these two tokens are at k tokens apart

• Below table defines a summary of APE vs RPE

• As we can generalise from above table that the two tokens which generally comes together or at k tokens apart, RPE will be same in such cases whereas in case of APE, it will differ for each sentence which model have to learn to generalise better
• So instead of saying i am at this position or that position, it would be better to explicitly mention that we are at k tokens apart

RoPE

• Let’s say we have a 2D vector and we want to rotate it by angle , where it was earlier at with magnitude r

• We can write initial and as:

  • =
  • =

• If we rotate it by , then:

  • = =
  • = =

• If we represent:

  

  

• Then we can write

• Below two properties which will be used later:

  • = which gives
  • = I [Identity Matrix]
  • =

• RoPE rotates the embeddings

• Let’s say q(query) at m and k(key) at n

• Then dot product between q and k would be

  • 
  • =
  • Since = , so =
  • this gives
  • now since =
  • this gives
  • In above eqn attention score now explicitly depends on relative distance between q and k
  • So if relative distance changes then only this factor will change

Implementation

• Treat a pair as single 2D vector space and rotate it independently
• For example lets say we have d_model of 256 then we will have 128 2D vectors and we can rotate those pairs independently as shown in below table

Pairs	Rotation Matrix
……………..	…………
……………..	…………

• where
  • is embedding value at index k
  • is rotation angle given as = , where i runs from 1 to
  • decreases as i increases, so initial embedding values are rotated by large angles as compared to later embedding values
• Below is the code to get and also validate if = I [Identity Matrix]
• For simplicity, i have just printed a single matrix

import torch

def get_RoPE(pos, d_model):
    theta = torch.tensor([10000**(-2*(i-1)/d_model) for i in range(1, d_model//2+1)])
    bsz, seq_len= pos.shape
    cos = torch.cos(pos[:, :, None]* theta[None, :]).view(bsz, seq_len, d_model//2)
    sin = torch.sin(pos[:, :, None]* theta[None, :]).view(bsz, seq_len, d_model//2)
    row1 = torch.stack([cos, -sin], dim=-1)
    row2 = torch.stack([sin, cos], dim=-1)

    RTheta = torch.stack([row1, row2], dim=-2)
    return RTheta

bsz = 2
tokens = torch.tensor([[2, 3] for _ in range(bsz)])
d_model = 256

RTheta = get_RoPE(tokens, d_model)

matrix = torch.matmul(RTheta, RTheta.transpose(-1, -2))

matrix[0][0][0]

tensor([[1.0000, 0.0000],
        [0.0000, 1.0000]])

• Below Code shows how to integrate RoPE while calculating attention

import torch
import torch.nn as nn

class SimpleAttention(nn.Module):
    def __init__(self, d_model, hidden_dim):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.q = nn.Linear(in_features = d_model, out_features = hidden_dim)
        self.k = nn.Linear(in_features = d_model, out_features = hidden_dim)
        self.v = nn.Linear(in_features = d_model, out_features = hidden_dim)
    
    def forward(self, x, rotary_emb):
        bsz, seq_len, d_model = x.shape
        query = self.q(x)
        key = self.k(x)
        value = self.v(x)

        if rotary_emb is not None:
            bsz, seq_len, d_model = x.shape
            query = query.view(bsz, seq_len, self.hidden_dim//2, 2, 1)
            key = key.view(bsz, seq_len, self.hidden_dim//2, 2, 1)
            query =  torch.matmul(rotary_emb, query).view(bsz, seq_len, self.hidden_dim)
            key =  torch.matmul(rotary_emb, key).view(bsz, seq_len, self.hidden_dim)
        
        scores = torch.softmax(torch.matmul(query, key.transpose(-1, -2))/(self.hidden_dim**2), dim=-1)
        final = torch.matmul(scores, value)
        return final

bsz = 2
tokens = torch.tensor([[2, 3] for _ in range(bsz)])
d_model = 256
embedding = torch.randn(*tokens.shape, d_model)

RTheta = get_RoPE(tokens, d_model)

model = SimpleAttention(d_model, d_model)

out = model(embedding, RTheta)
print(out.shape)

torch.Size([2, 2, 256])

Positional Encoding (PE)

• RNNs process tokens sequentially so they have information about the position of tokens of a given sequence.
• Transformers process all the tokens simultaneously where each token attends to every other token.
• For Example::-
  • my daughter called her brother [where each word is a token]
  • in above sentence my will attend to all other tokens and similarly others will also attend other tokens.
  • now consider other permutation of above sequence her daughter called my brother
  • even in above example, permutation is different but transformers will find out the same context for each word as above one
  • a simple explanation would be consider my and her
    • the attention between these will be same for both of the sequences, but both of them corresponds to different meaning.
    • so to have different attentions for different sequences we need to inject some information about the positions of words also.

Attempted Solution

• Use the index of tokens directly as position information for a given sequence. However, when we add token embedding with this position indices, positional information will overpower the content signal and also, the rightmost tokens will dominate.

• Now, to solve this, we can normalize the positional information i.e. normalize with the length of sequence. But now we have another problem, all of the sequences will never be of same length, so each position information will keep varying. For example lets consider two sequences

  1. I am going PE:=[0/3 1/3 2/3]
  2. All of us are going PE:=[0/5 1/5 2/5]

• As we can see in above example, in sequence we are normalizing by 3 and in 2nd one by 5, which is not good for a model to learn like one time we are saying that 2nd positions value is 0.33 and in the next sequence we are saying no no, its 0.2

Sinusoidal Positional Encoding

• The authors of the Transformer paper introduced sinusoidal functions to address this issue.

• The positional encoding is defined as:

  • For even indices:

  

  • For odd indices:

• where
  • pos is the position of token
  • is dimensionality of position embedding
  • i is index of embedding dimension where i runs from 0 to /2 -1 [each value of i corresponds to one frequency of sinusoid]

import torch
from typing import Union

def get_1d_sin_cos_pe(d_model:int, pos:Union[int, list]):
    if type(pos)==int:
        pos = [pos]
    n_tokens = len(pos)
    pos = torch.tensor(pos).reshape(n_tokens, -1)
    div_term = torch.tensor([10000**(2*i/d_model) for i in range(d_model//2)])
    sinPE = torch.sin(pos/div_term)
    cosPE = torch.cos(pos/div_term)
    combinedPE = torch.zeros(n_tokens, d_model)
    combinedPE[:, 0::2] = sinPE
    combinedPE[:, 1::2] = cosPE

    return combinedPE

# For single position
d_model = 512
pos = 2
pe2 = get_1d_sin_cos_pe(d_model, pos)

d_model = 512
pos = 3
pe3 = get_1d_sin_cos_pe(d_model, pos)

combinedPE = torch.cat([pe2, pe3], dim=0)

# For multiple positions 
d_model = 512
pos = [2, 3]
pe = get_1d_sin_cos_pe(d_model, pos)

print(torch.equal(combinedPE, pe))

True

Lets see only sine for intuition

• As the scaling factor increases, the period of sin wave increases frequency decreases i.e. the repetition of values decreases.

• We can also say that as scale factor increases, the oscillation becomes slower.

• Dimension 0 i.e. lower dimension can incorporate fine grained local positional information.

• Value at k is far different than value at k + 5 as we can see from below figure.

---

• Dimension 255 i.e. higher dimensions incorporates global positional information
• Value at k and k+5 is nearly similar but at k and k+200 is quite different as we can see from the below figure.

---

• Basically given a particular position pos, it gets mapped to various sine waves which incorporates local as well as global information as we can see from below figure.

Why 10000?

What happens if we chose 100

• From below graph we can see that even higher dimension starts repeating themselves i.e. their frequency increases and thus we cant distinguish between far-apart tokens

• For example lets calculate cosine similarity between PE[128:256] of 500th position to everyone other position.
• The cosine similarity will change drastically as we can see in the below figure, whereas we want higher dimensions to change smoothly.

• The cosine similarity between PE[0:128] of 500th position to everyone other position.
• The cosine similarity decreases drastically to zero even with position difference of 10 only, whereas we want it to decrease smoothly but not in an aggressive way.

What happens if we chose 100000

• And if we chose 100K as scaling factor then from below figure we can see that even the earlier dimensions have reduced frequency and thus we can’t differentiate between two local positions

• For example lets calculate cosine similarity between PE[128:256] of 500th position to everyone other position, in this case similarity barely changes as we can see from below figure.

• The cosine similarity between PE[0:128] of 500th position to everyone other position.
• The cosine similarity remains mostly same even with position difference of 50, whereas we want it to decrease smoothly.

How does 10k solves above problems

• And as suggested in paper, if we chose 10k. It is something like sweet spot where earlier dimensions have higher frequency thus helping in distinguishing between local positions and later dimensions have lower frequency thus helping in capturing long range positional differences.

• As we can see from below graphs it solves the problem of cosine similarity between PE of 500th position and every other position

Why include cosine also

• It allows, the PE across all positions to have simple linear relationship.
• Generally:
  • x =
  • y =
• But in paper:
  • x =
  • y =

• As we can see from above equations, PE[pos+k] = PE[pos]

Chapter 1 Introduction to Building AI Applications with Foundation Models

Chapter 2 Understanding Foundation Models

Transformer Architecture

Transformer Block

Each transformer block consists of: - Attention module - MLP (feedforward network)

Transformer-Based Language Model

A typical transformer-based LM includes:

• Embedding Module (pre-transformer)
  • Converts tokens into embedding vectors
• Transformer Blocks
• Model Head (post-transformer)
  • Converts hidden states into token probabilities

---

Pre-training and Post-training

• Pre-training
  • Equivalent to reading to acquire knowledge
  • Resource-intensive
  • Produces general-purpose representations
• Post-training
  • Equivalent to learning how to use knowledge
  • Includes instruction tuning, alignment, and RLHF

• Training only on high-quality data without pre-training is possible, but pre-training followed by post-training yields superior results

---

Chapter 3 Evaluation Methodology

Motivation

Scaled Dot-Product Attention

PyTorch Implementations

import torch
import torch.nn as nn
import torch.nn.functional as F
import time

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class SingleHeadVanillaAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.d_model = d_model
        self.d_k = d_k
        self.Wq = nn.Linear(in_features = self.d_model, out_features = self.d_k)
        self.Wk = nn.Linear(in_features = self.d_model, out_features = self.d_k)
        self.Wv = nn.Linear(in_features = self.d_model, out_features = self.d_k)

    def forward(self, x):
        query = self.Wq(x) # [B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)
        key = self.Wk(x)   #[B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)
        value = self.Wv(x) # [B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)

        scores = torch.matmul(query, key.transpose(-1, -2))/(self.d_k**0.5) # [B, N, d_k] * [B, d_k, N] -> [B, N, N] TC=O(B*N^2*d_k)

        attn = F.softmax(scores, dim=-1) # [B, N, N]

        final = torch.matmul(attn, value) # [B, N, N] * [B, N, d_k] -> [B, N, d_k]  TC=O(B*N^2*d_k)
        return final # [B, N, d_k]

d_model = 512
d_k = 512

model = SingleHeadVanillaAttention(d_model = d_model, d_k = d_k).requires_grad_(False).eval().to(device)

n_tokens = 16384
inp = torch.randn(1, n_tokens, d_model, device=device)

torch.cuda.synchronize()
start = time.time()
k=50
for _ in range(k):
    with torch.no_grad():
        out = model(inp)
torch.cuda.synchronize()
print(f"Time taken: {(time.time()-start)/k:.3f}")

Time taken: 0.508

import torch
import torch.nn as nn
import torch.nn.functional as F
import time

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class MultiHeadVanillaAttention(nn.Module):
    def __init__(self, d_model, num_heads, head_dim=None):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        if head_dim is None:
            self.d_k = d_model//self.num_heads
        else:
            self.d_k = head_dim
        self.Wq = nn.ModuleList([nn.Linear(in_features = self.d_model, out_features = self.d_k) for _ in range(self.num_heads)])
        self.Wk = nn.ModuleList([nn.Linear(in_features = self.d_model, out_features = self.d_k) for _ in range(self.num_heads)])
        self.Wv = nn.ModuleList([nn.Linear(in_features = self.d_model, out_features = self.d_k) for _ in range(self.num_heads)])
        self.out = nn.Linear(in_features = self.d_k * self.num_heads, out_features = self.d_model)

    def forward(self, x):
        query = [self.Wq[i](x) for i in range(self.num_heads)] # [B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)
        key = [self.Wk[i](x) for i in range(self.num_heads)]   #[B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)
        value = [self.Wv[i](x) for i in range(self.num_heads)] # [B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)

        scores = [torch.matmul(query[i], key[i].transpose(-1, -2))/(self.d_k**0.5) for i in range(self.num_heads)] # [B, N, d_k] * [B, d_k, N] -> [B, N, N] TC=O(B*N^2*d_k)

        attn = [F.softmax(scores[i], dim=-1) for i in range(self.num_heads)] # [B, N, N]

        final = [torch.matmul(attn[i], value[i]) for i in range(self.num_heads)] # [B, N, N] * [B, N, d_k] -> [B, N, d_k]  TC=O(B*N^2*d_k)
        final = torch.cat(final, dim=-1)

        projout = self.out(final)

        return projout # [B, N, d_model]

d_model = 512
d_k = 512

model = MultiHeadVanillaAttention(d_model = d_model, num_heads=2, head_dim=32).requires_grad_(False).eval().to(device)

n_tokens = 16384
inp = torch.randn(1, n_tokens, d_model, device=device)

torch.cuda.synchronize()
start = time.time()
k=50
for _ in range(k):
    with torch.no_grad():
        out = model(inp)
torch.cuda.synchronize()
print(f"Time taken: {(time.time()-start)/k:.3f}")

Time taken: 0.926

import torch
import torch.nn as nn
import torch.nn.functional as F
import time

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class MultiHeadVanillaAttention(nn.Module):
    def __init__(self, d_model, head_dim, num_heads):
        super().__init__()
        self.d_model = d_model
        self.d_k = head_dim * num_heads
        self.head_dim = head_dim
        self.num_heads = num_heads
        self.Wq = nn.Linear(in_features = self.d_model, out_features = self.d_k)
        self.Wk = nn.Linear(in_features = self.d_model, out_features = self.d_k)
        self.Wv = nn.Linear(in_features = self.d_model, out_features = self.d_k)

        self.out = nn.Linear(in_features = self.d_k, out_features = self.d_model)
    def forward(self, x):
        batch_size = x.shape[0]
        query = self.Wq(x) # [B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)
        key = self.Wk(x)   #[B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)
        value = self.Wv(x) # [B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)

        query = query.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) # [B, N, d_k] -> [B, N, num_heads, head_dim] -> [B, num_heads, N, head_dim]
        key = key.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) # [B, N, d_k] -> [B, N, num_heads, head_dim] -> [B, num_heads, N, head_dim]
        value = value.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) # [B, N, d_k] -> [B, N, num_heads, head_dim] -> [B, num_heads, N, head_dim]


        scores = torch.matmul(query, key.transpose(-1, -2))/(self.head_dim**0.5) # [B, num_heads, N, head_dim] * [B, num_heads, head_dim, N] -> [B, num_heads, N, N] TC=O(B*N^2*d_k)

        attn = F.softmax(scores, dim=-1) # [B, num_heads, N, N]

        final = torch.matmul(attn, value).permute(0, 2, 1, 3).reshape(batch_size, -1, self.d_k) # [B, num_heads, N, N] * [B, num_heads, N, head_dim] -> [B, N, d_k]  TC=O(B*N^2*d_k)
        projout = self.out(final) # [B, N, d_k] -> [B, N, d_model]
        return projout # [B, N, d_model]


d_model = 512
d_k = 512

model = MultiHeadVanillaAttention(d_model = d_model, num_heads=2, head_dim=32).requires_grad_(False).eval().to(device)

n_tokens = 16384
inp = torch.randn(1, n_tokens, d_model, device=device)

torch.cuda.synchronize()
start = time.time()
k=50
for _ in range(k):
    with torch.no_grad():
        out = model(inp)
torch.cuda.synchronize()
print(f"Time taken: {(time.time()-start)/k:.3f}")

Time taken: 2.734

Linear Attention

• Instead of using softmax, it tries to make it more general where this softmax function can be replaced by any other function

• Kernel is some function which takes in two input vectors (query and key) and gives a single scalar output i.e. similarity score.

• This softmax can also be replaced by some kernel which is some kind of similarity function (i.e. to have property of a similarity function).

• Kernel can be decomposed into linear functions (feature functions)

  • 

• Feature function takes a vector as input and projects it into new feature space

• The basic idea is instead of doing complicated non-linear function like softmax, can’t we just project a and b into higher dimensional space and just do the linear inner product

• The attention mechanism is defined as .

• We can make the above equation more generalized by using a similarity function

• 

• where we can think of

• only constraint on sim is to give non-negative outputs

• given a kernel with feature representation , we can rewrite the above equation

• 

• By making the use of associative property of matrix multiplication, we can take out Q from summation, because it is independent of j

• (4)

• And finally we can convert it into vectorized format

• 

• It is linear because we can compute numerator and denominator and use it for all the queries [The eqn before vectorized one]

• For softmax attention, complexity is O(N^2 * max(D, M)), where D is dimensionality for query and key, M is for value

• For linear attention, complexity is O(NCM), where C is feature map projection dim and M is for value

• Linearization of exact softmax attention is not feasible because of the exponential kernel, which is of infinite dimension and can’t be stored on a computer (expansion of e^x is infinitely many terms)

• That’s why it is better to use some kind of polynomial feature representation

• The author chose elu instead of relu to avoid setting gradients to 0 when input was negative

Single Head Linear Attention PyTorch Code

import torch
import torch.nn as nn
import torch.nn.functional as F
import time

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class SingleHeadLinearAttention(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.Wq = nn.Linear(in_features = self.in_features, out_features = self.out_features)
        self.Wk = nn.Linear(in_features = self.in_features, out_features = self.out_features)
        self.Wv = nn.Linear(in_features = self.in_features, out_features = self.out_features)

    def forward(self, x):
        query = self.Wq(x) # [B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)
        key = self.Wk(x)   #[B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)
        value = self.Wv(x) # [B, N, d_model] * [d_model, d_k] -> [B, N, d_k] TC=O(B*N*d_model*d_k)

        query = F.elu(query) + 1 # [B, N, d_k]
        key = F.elu(key) + 1 # [B, N, d_k]

        scores = torch.matmul(key.transpose(-1, -2), value) # [B, d_k, N] * [B, N, d_k] -> [B, d_k, d_k] TC=O(B*N*d_k^2)

        final = torch.matmul(query, scores) # [B, N, d_k] * [B, d_k, d_k] -> [B, N, d_k] TC=O(B*N*d_k^2)
        k_sum = torch.sum(key, dim=-2).unsqueeze(-1)

        normalizer = torch.matmul(query, k_sum) # [B, N , d_k] * [B, d_k, 1] -> [B, N, 1] TC=O(B*N*d_k)
        eps = 1e-6 # To use in case we don't divide by 0

        return final/(normalizer + eps) # [B, N, d_k]

d_model = 512
d_k = 512

model = SingleHeadLinearAttention(in_features = d_model, out_features = d_k).requires_grad_(False).eval().to(device)

n_tokens = 16384
inp = torch.randn(1, n_tokens, d_model, device=device)

torch.cuda.synchronize()
start = time.time()
k=50
for _ in range(k):
    with torch.no_grad():
        out = model(inp)
torch.cuda.synchronize()
print(f"Time taken: {(time.time()-start)/k:.3f}")

Time taken: 0.057

Method	Main idea	Complexity bottleneck
Single Head Vanilla Attention	One set of query, key, and value projections	Forms an attention matrix
Multi Head Vanilla Attention For loop	Computes each head separately and concatenates outputs	Still forms one matrix per head
Multi Head Vanilla Attention	Vectorizes the heads into one batched operation	Same quadratic attention cost, but better parallelism
Single Head Linear Attention	Uses a feature map and associativity to avoid explicit softmax attention	Depends on feature dimension instead of storing full attention

Chapter-1[Introduction to LLMs]

• Bag of Words:- We have a predefined vocabulary, which is used to creat a vector, where each vector index contains the count of that particular word index in vocab

• word2vec:- It is trained by using a pair(whether dfferent or same semantic meaning) of words

• Bag of words create embedding at documet level, whereas word2vec generates embedding for words only

• Representation Models are encoder only models

• Generative Models are decoder only models

• BERT(Bidirectional Encoder Represenations from Transformers)

  • It is trained by a tecique called masked language modeling, where this method masks a part of the input for the model to predict
  • Pretrain on large dataset using masked language modng and then fine tune it for down stream tasks

• Creating an LLM consists of typically at least two steps:

  • Language Modeling(Pretraining):- An LLM is trained on a vast corpus of internet text allowing the model to learn grammar, context and language patterns.
  • Fine Tuning(Post Training):- Further training the previously trained model on narrower tasks.

Chapter-2[Tokens and Embeddings]

• Partial tokens (like “izing”, and “ic”) have a special hidden character at their beginning that indicates that they are connected with the token that precedes them in the text.

• Word Tokens

• Subword Tokens

• Character Tokens

• Byte Tokens

• Designing Large Language Model Applications. (For tokenizers)

• BERT tokenizers are based on wordpiece, introduced in Japanese and Korean voice search [https://ieeexplore.ieee.org/document/6289079]

• GPT-2 is based on Byte Pair Encoding (BPE), introduced in Neural machine translation of rare words with subword units [https://arxiv.org/abs/1508.07909]

• Flan-T5 is based on SentencePiece, introduced in “SentencePiece: simple and language independent subword tokenizer and detokenizer for neural text processing”, which supports BPE and unigram language model

• There are three major factors that dictate how a tokenizer breaks down an input prompt

  • Tokenization method (BPE, WordPiece, etc.)
    • Each of these methods outlines an algorithm for how to choose an appropriate set of tokens to represent a dataset
  • Tokenizer Design choices (vocab size, special tokens)
    • vocab size:- How many tokens to keep in tokenizers vocabulary
    • special tokens:- What special tokens do we want the model to keep track of. We can add as many of these as we want
  • Capitalization:- Whether to convert caps to lower or not

• More details on training tokenizer at:-

  • Tokenizers section of the Hugging Face course [https://huggingface.co/learn/nlp-course/chapter6/1?fw=pt]
  • Natural Language Processing with Transformers, Revised Edition [https://www.oreilly.com/library/view/natural-language-processing/9781098136789/]

Chapter-3[Looking Inside Large Language Models]

• Autoregressive Models: The models that consume their earlier predictions to make later predictions

• Three Major components of LLMs are:

  • Tokenizer
  • Stack of Transformer Blocks
  • A language modeling head

• Context length (Number of streams):- Number of previous tokens which will be considered for the prediction of current token

• KVCache:- Keeping the previously calculated key and value so that we dont have to recalculate the same again and again

• Components of Transformer Blocks

  • Attention Layer:- Incorporates contextual information to allow the model to better capture the nuance of language.
  • FeedForward Layer:- It is able to store information and make predictions and interpolations from data it was trained on.

• Two main steps are involved in the attention mechanism:

  • Relevance scoring: How much previous tokens are relevant to the curren token being processed.
  • Combining Information:- Using scores, combine the information from the various positions into a single output vector

• More efficient Attention

  • Local/Sparse Attention:- Sparse Attention limits the context of previous tokens that the model can attend to.
  • Multi-query:- Unique Queries but same key and value for all the head
  • Grouped-query:- Dividing queries in groups and then key, value will be same for the queries within a group
  • RoPE:- Added to queries and keys before calculating the score in Attention blocks

Chapter-4[Text Classification]

• In case of no lableld data, we can define our desired labels and then encode both labels and given text and then use cosine_similarity

• Directly using pretrained models for sentiment classification

• Using a simple classifier on top of embedding generator

• If we don’t have labeled data then we can use cosine similarity to find out the label

• We can use generative models also

Chapter-5[Text Clustering and Topic Modeling]

• A common Pipeline for Text Clustering

  • Convert the input documents to embeddings with an embedding model
  • Reduce the dimensionality of embeddings with a dimensionality reduction model (UMAP as it tends to handle nonlinear relationships and structures a bit better than PCA)
  • Find groups of semantically similar documents with a cluster model (using HDBSCAN)

• c-TF-IDF

  • TF:- Frequency of word X in class C
  • IDF:- log(average number of words per class/frequency of X across all classes)

• BERTopic [https://maartengr.github.io/BERTopic/getting_started/best_practices/best_practices.html]

• After getting the keywords using c-TF-IDF, we can use MMR(maximal marginal relevance) to keep only the diverse keywords, we can also use KeyBERTInspired to fine-tune the topic representations

• Additonally, we can use LLMs to further improve interpretability of topics

Chapter-6[Prompt Engineering]

• Temperature:- It controls the randomness or creativity of the text generated. It defines how likely it is to choose tokens that are less probable. A temperature of 0 generates the same response every time because it always choose the most likely word. A higher value allows less probable words to be generated

• Divide all the logits before passing to softmax

• top_p:- It is also known as nucleus sampling, is a sampling technique that controls which subset of tokens (the nucleus) the LLM can consider. It will consider tokens until it reaches their cumulative probability. If we set top_p to 0.1, it will consider tokens until it reache that value.

• top_k:- same as top_p, LLM will only consider the top 100 most probable tokens if you set its value to 100.

• Table 6.1 from page 172

• Self-consistency:- Sampling from multiple paths for the same prompt and using the voting technique for the best one

• CoT:- Solving the complex problems step by step

• ToT:- Generates different solutions and then selects the best one and continues. This method requires so many calls to the model. But we can change it using a prompt, we can ask the model to mimic the behavior

Chapter-7[Advanced Text Generation Techniques and Tools]

• PromptTemplate from langchain
• Can use LLMChain to create a chain
• Can use ConversationBufferMemory to have access to chat history
• Can use ConversationBufferWindowMemory to keep only last k chats–> Can use ConversationSummaryMemory alongside an LLM to store the summary fo conversation instead of storing raw chats
• ReAct:- Reasoning and Acting
  • Thought:- thought about the input prompt
  • Action:- based on thought action is triggered, it is generally an external tool like calculator or search engine
  • Observation:- finally the results of action are returned to the LLM it observes the output

Chapter-8[Semantic Search and Retrieval-Augmented Generation]

Theory vs practice:-

• If you aim to reach the highest echelons of deep learning research, you first and foremost need to be a great practitioner!
• For best effects, use one cup of theory, one cup of practice. Rinse and repeat
• We are reading papers, but without training actual models, without the experience of applying what we are learning to real-life problem, we are missing a very important feedback loop.

Programming is about what you have to say:-

• Your ability as a developer is measured by the utility of the things you can express in a language.
• Reading and writing a lot of code is the best way to learn a programming language.

The secret of good developers:-

• Go for long, uninterrupted sessions to avoid context switching.

The best way to improve as a developer:-

• You don’t sharpen your skills with resources, books, or articles. You sharpen your skills with practice. If you want to get better, go do the thing.
• Read and write a lot of code.

How to use your tools to achieve a state of flow

• You are in a state of flow when nothing separates you and your work
• That’s why many programmers spend a lot of time studying their editors
• We are always somewhere between two extremes (state of flow)
• All we can do is take small steps and see if it is taking us towards the ideal.
• This can take our productivity to a whole new level.

Use reality as your mirror:-

• If you want to live a different life tomorrow than you are living today, you have to put your beliefs to the test.
• The stronger the emotions involved in a situation, the higher the chance you are not seeing things clearly.
• Suppose you want to post something and you are afraid that people whom I respect will unfollow me or they will say I am stupid, but that doesn’t matter in the long run.
• Post it and you will learn on the way.
• Thinking about writing something isn’t doing the thing.

Do genuine work (it compounds):-

• Reading a book without taking notes is like discovering a new territory and forgetting to draw a map.
• Genuine work moves the atoms in the universe.
• Always make notes.
• Share whatever you have made.
• The more atoms you move the better the feedback you receive and the greater the extent to which you can reflect on what you are learning.

The hidden game of machine learning:-

• At the core of machine learning lies the ability to generalize to unseen data.
• How (and why) to create a good validation set blog post by Rachel Thomas
• Learning From Data: A short course by Yaser S. Abu-Mostafa
• Building Machine Learning Powered Applications: Going from Idea to Product by Emmanuel Ameisen

How to structure a machine learning project:-

• The main condition for a healthy machine learning project is a good train-val-test split.

• Always have a baseline to get to know whether i am going in right direction or not or what can be done next.

• Don’t run experiments just for the sake of tweaking hyperparameters, instead invest time in exploring other architectures, developing diagnostic code.

• The idea is to always move in small increments, using simpler models as a stepping stone to more complex ones.

• Define loss and implement it on simple model consisting of few layers and then check by passing a batch if output is not all zeros and the output shape is required shape.

• Create a subset of train data(1%) to check whether the modified code is working and the we are going in the right direction.

• Read as much as possible related to the domain you are working on so that our mind can work on the solution and analyze what to do next.

• Go slowly and give attention to each component of the pipeline and grow them systematically.

How to win at Kaggle:-

• Join the competition early.

• Learn about the problem domain and data

• Write diagnostic code and identify where model is struggling.

• Follow Kaggle forums and read daily.

• Try to find out the relevant research papers and see if it is useful in solving the problem.

• You can read blogs related to the problem you are solving.

• While early in the competition, try on different architectures instead of tweaking the hyperparameters.

• How do we know that we are moving in the right direction? So first of all create a baseline and make a submission. The submission can consist of all zeros.

• Try to find out a validation split that will track the public leaderboard.

• A proper validation split lies at the heart of any good Kaggle submission.

• To know that if we are going in the right direction, see if the local results matches with the leaderboard. If so then you are going in the right direction and if not then try to find out other ways of splitting the data or something else.

• You can do ensembling also. the idea is to cancel out the errors from individual model in a hope that their errors are not correlated.

• A related technique is training with cross-validation, where you train multiple models withholding different parts of the data for validation and then combine the results from them.

The best hardware for deep learning:-

• Follow this blog to know more about GPUs

Debugging with ease is a superpower:-

• Radek Osmulski’s Favorite recent Jupyter notebook discovery - the %debug magic:

Time yourself:-

• Always keep track of the time which you devote to various activities.
• %%timeit

You can’t learn a profession by studying a textbook:-

• Extensive practice is very important alongside theory.

On finding a job:-

• The best approach is to showcase your work, by helping others.
• Reach out to the people who know you and can help you.

The deep learning party is on Twitter:-

• you will find all the amazing people there

Share your work:-

• Making Peace with Personal Branding

When to start sharing your work:-

• The sooner, the better!
• With every piece that we produce, we hone our communication skills

I am scared to share my work! Help!

• Where some see failure, others see cheap feedback.
• If you don’t like something then you can always delete it.

What to focus on in sharing your work:-

• Speak to your experience.
• You don’t have to share something outside of your experience.

Don’t lose sight of what is important:-

• You won a Kaggle competition that’s good but never ever forget why you have started.
• You don’t need followers or likes or even social media to learn.
• Nearly all of the best DL engineers I(Radek Osmulski) know are either very quiet or completely silent on Twitter

Make mental space for what matters:-

• Peace of mind is the most important prerequisite for creative work.
• Avoid using social media too much!

To engage afterburners, find a mentor:-

• Find a mentor, who is good at something you care about.
• The mentor doesn’t need to know you.
• In general, we want to get answers from the people who are like us and have gone through that path and the people who are not like us to get to know about more ideas.
• Always make your message as concise and clear as you can.

The biggest regret of fast.ai students:-

• I wished I spent more time coding and experimenting and less time studying in a more traditional sense.
• 100% learning -> x% doing/100-x% real learning.
• Shortest path to understand how something works leads through practice.
• 80% doing and 20% reading theory.

Persistence is everything:-

• At any given moment, as you put in the work, you can barely notice a difference in your life. But the longer you stay the course the more rewarding the journey becomes.
• Learning compounds and you need to give it time before you start seeing the exponential results.
• Combine persistence with community involvement and you cannot be stopped.

Change is about what not to do:-

• Ceasing how we have done things up to now is harder than taking on new approach.
• Try to find out the behaviors which are no longer helpful to pursuing your goals and replace them with some meaningful.

Learning might just be enough

• A good strategy towards learning is to observe whether you are getting the results that you are after and if you are not then change your approach.

More perspectives on mentoring:-

• If you want to catch the attention of someone you admire give yourself a job working for them.
• Learn in public.
  • Write blogs and tutorials and cheatsheets.
  • Speak at meetups and conferences.
  • Ask and answer things on Stack Overflow or Reddit. Avoid the walled gardens like Slack and Discord, they’re not public.
  • Make YouTube videos or Twitch streams.
  • Start a newsletter.
  • Draw cartoons
• Make the thing you wish you had found when you were learning.
• Don’t judge results by claps or tweets or stars or upvotes, just talk to yourself from 3 months ago.
• by far the biggest beneficiary of you trying to help past you is future you. If others benefit, that’s icing.
• Don’t stop there:
  • Enjoyed a coding video? Reach out to the speaker/instructor and thank them, and ask questions.
  • Make PR’s to libraries you use.
  • Make your own libraries no one will ever use.
  • Clone stuff you like, from scratch, to see how they work.
  • Teach workshops.
  • Go to conferences and summarize what you learned.
• Try your best to be right, but don’t worry when you’re wrong, and let the internet correct you when you are inevitably wrong.
• Talk while you code.
• Whenever someone wants some help, then you should be always ready to help them because these are some of the most in-demand people in tech.

Tap into the power of the community to learn faster

• Instead of asking questions only try to answer more.
• Joining community will help you in better communication, learn from other’s.

Energize

• Yoga can help.
• Time-restricted eating and its primary benefit is restore hormonal balance.

Quantize and De-quantize a tensor

• Advantages of Quantization
  • Smaller model
  • Speed gains
    • Memory bandwidth
    • Faster operations
      • GEMM: General Matrix Multiply(matrix to matrix multiplication)
      • GEMV: General Matrix Vector Multiplication (matrix to vector multiplication)
• Challenges of Quantization
  • Quantization error
  • Retraining (Quantization Aware Training)
  • Limited Hardware support
  • Calibration dataset needed
  • packing/unpacking
• getting q:-
  • r = s(q-z) q = int(round(r/s+z))

import torch
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

#Helper functions to visualize
def plot_matrix(tensor, ax, title, vmin=0, vmax=1, cmap=None):
    """
    Plot a heatmap of tensors using seaborn
    """
    sns.heatmap(tensor.cpu().numpy(), ax=ax, vmin=vmin, vmax=vmax, cmap=cmap, annot=True, fmt=".2f", cbar=False)
    ax.set_title(title)
    ax.set_yticklabels([])
    ax.set_xticklabels([])


def plot_quantization_errors(original_tensor, quantized_tensor, dequantized_tensor, dtype = torch.int8, n_bits = 8):
    """
    A method that plots 4 matrices, the original tensor, the quantized tensor
    the de-quantized tensor and the error tensor.
    """
    # Get a figure of 4 plots
    fig, axes = plt.subplots(1, 4, figsize=(15, 4))

    # Plot the first matrix
    plot_matrix(original_tensor, axes[0], 'Original Tensor', cmap=ListedColormap(['white']))

    # Get the quantization range and plot the quantized tensor
    q_min, q_max = torch.iinfo(dtype).min, torch.iinfo(dtype).max
    plot_matrix(quantized_tensor, axes[1], f'{n_bits}-bit Linear Quantized Tensor', vmin=q_min, vmax=q_max, cmap='coolwarm')

    # Plot the de-quantized tensors
    plot_matrix(dequantized_tensor, axes[2], 'Dequantized Tensor', cmap='coolwarm')

    # Get the quantization errors
    q_error_tensor = abs(original_tensor - dequantized_tensor)
    plot_matrix(q_error_tensor, axes[3], 'Quantization Error Tensor', cmap=ListedColormap(['white']))

    fig.tight_layout()
    plt.show()

def linear_q_with_scale_and_zero_point(
    tensor, scale, zero_point, dtype = torch.int8):

    scaled_and_shifted_tensor = tensor / scale + zero_point

    rounded_tensor = torch.round(scaled_and_shifted_tensor)

    q_min = torch.iinfo(dtype).min
    q_max = torch.iinfo(dtype).max

    q_tensor = rounded_tensor.clamp(q_min,q_max).to(dtype)
    
    return q_tensor

test_tensor=torch.tensor(
    [[191.6, -13.5, 728.6],
     [92.14, 295.5,  -184],
     [0,     684.6, 245.5]])

scale = 3.5
zero_point = -70

quantized_tensor = linear_q_with_scale_and_zero_point(
    test_tensor, scale, zero_point)

def linear_dequantization(quantized_tensor, scale, zero_point):
    return scale * (quantized_tensor.float() - zero_point)

dequantized_tensor = linear_dequantization(
    quantized_tensor, scale, zero_point)

plot_quantization_errors(test_tensor, quantized_tensor,
                         dequantized_tensor)

Get the Scale and Zero-Point

• s = (r_max-r_min)[current_tensor_range]/(q_max-q_min)[datatype_range]
• z = int(round(q_min - r_min/s))
• z and quantized tensor are of the same type
• z is an integer because it represent zero(in the original ‘r’ range) with an integer in the quantized ‘q’ range
• if z goes out of range:-
  • z < q_min:-
    • z = q_min
  • z > q_max:-
    • z = q_max

import torch

def get_q_scale_and_zero_point(tensor, dtype=torch.int8):
    
    q_min, q_max = torch.iinfo(dtype).min, torch.iinfo(dtype).max
    r_min, r_max = tensor.min().item(), tensor.max().item()

    scale = (r_max - r_min) / (q_max - q_min)

    zero_point = q_min - (r_min / scale)

    # clip the zero_point to fall in [quantized_min, quantized_max]
    if zero_point < q_min:
        zero_point = q_min
    elif zero_point > q_max:
        zero_point = q_max
    else:
        # round and cast to int
        zero_point = int(round(zero_point))
    
    return scale, zero_point

Symmetric vs Asymmetric Mode

• Asymmetric Mode:-
  • map [r_max, r_min] to [q_max, q_min]
  • This is what we have implemented above
• Symmetric Mode:-
  • map [-r_max, r_max] to [-q_max, q_max]
    • where r_max = max(|tensor|)

Hence, we can simplify the equation to:-

• q = int(round(r/s))
• s = r_max/q_max

import torch

def get_q_scale_symmetric(tensor, dtype=torch.int8):
    r_max = tensor.abs().max().item()
    q_max = torch.iinfo(dtype).max

    # return the scale
    return r_max/q_max

def linear_q_symmetric(tensor, dtype=torch.int8):
    scale = get_q_scale_symmetric(tensor)
    
    quantized_tensor = linear_q_with_scale_and_zero_point(tensor,
                                                     scale=scale,
                   # in symmetric quantization zero point is = 0    
                                                    zero_point=0,
                                                      dtype=dtype)
    
    return quantized_tensor, scale

Trade-off

• Utilization of quantized range:

  • when using asymmetric quantization, the quantized range is fully utilized
  • When symmetric mode, if the float range is biased towards one side, this will result in a quantized range where a part of the range is dedicated to values that we’ll never see.(e.g ReLU where the output is positive)

• Simplicity:

  • Symmetric mode is much simpler compared to asymmetric mode.

• Memory: We don’t have to store zero-point for symmetric quantization

• We use symmetric quantization for 8-bit, but as we go for lower bits such as 2 or 4 bits, we use asymmetric quantization

Finer Granularity for more Precision

• Different granularities
  • per tensor
  • per channel (along an axis)
  • per group (group n elements together)
• The more granular quantization is the more precise it will be.

Per Channel Quantization

• we usually use per channel quantization in int8

def linear_q_symmetric_per_channel(r_tensor, dim, dtype=torch.int8):
    
    output_dim = r_tensor.shape[dim]
    # store the scales
    scale = torch.zeros(output_dim)

    for index in range(output_dim):
        sub_tensor = r_tensor.select(dim, index)
        scale[index] = get_q_scale_symmetric(sub_tensor, dtype=dtype)

    # reshape the scale
    scale_shape = [1] * r_tensor.dim()
    scale_shape[dim] = -1
    scale = scale.view(scale_shape)
    quantized_tensor = linear_q_with_scale_and_zero_point(
        r_tensor, scale=scale, zero_point=0, dtype=dtype)
   
    return quantized_tensor, scale

Per Group Quantization

• Group n(e.g. 32, 64, 128) elements together and quantize

• Per group quantization can require a lot of memory

  • Let’s say we want to quantize a tensor in 4-bit and we choose group_size=32, symmetric mode(z=0), and we store the scales in FP16
  • It means that we actually quantizing the tensor in 4.5 bits since we have:
    • 4 bit(each element is stored in 4 bit)
    • 16/32 bit (scale in 16 bits for every 32 elements)

def linear_q_symmetric_per_group(tensor, group_size,
                                 dtype=torch.int8):
    
    t_shape = tensor.shape
    assert t_shape[1] % group_size == 0
    assert tensor.dim() == 2
    
    tensor = tensor.view(-1, group_size)
    
    quantized_tensor, scale = linear_q_symmetric_per_channel(
                                tensor, dim=0, dtype=dtype)
    
    quantized_tensor = quantized_tensor.view(t_shape)
    
    return quantized_tensor, scale

def linear_dequantization_per_group(quantized_tensor, scale, 
                                    group_size):
    
    q_shape = quantized_tensor.shape
    quantized_tensor = quantized_tensor.view(-1, group_size)
    
    dequantized_tensor = linear_dequantization(quantized_tensor, 
                                               scale, 0)
    
    dequantized_tensor = dequantized_tensor.view(q_shape)
    
    return dequantized_tensor

Quantizing Weights and Activations for Inference

• Depending on what we quantize, the storage and the computation are not the same.
• W8A32
  • If weights are quantized but not the activations, then computation is done floating point (FP16,FP32, BF16)
  • We need to dequantize the weights to perform the floating point computation (cast to float32)
• W8A8
  • Both are quantized
  • Computation is integer based but not supported by all hardware

Custom Build an 8-Bit Quantizer

#W8A16LinearLayer
def w8_a16_forward(weight, input, scales, bias=None):
    
    casted_weights = weight.to(input.dtype)
    output = F.linear(input, casted_weights) * scales
    
    if bias is not None:
        output = output + bias
      
    return output

import torch
import torch.nn as nn
import torch.nn.functional as F

class W8A16LinearLayer(nn.Module):
    def __init__(self, in_features, out_features, 
                 bias=True, dtype=torch.float32):
        super().__init__()
        
        self.int8_weights = nn.Parameter(torch.Tensor([0, 1]
                                     ).to(dtype=torch.int8))

try:
    
    W8A16LinearLayer(1, 1)
    
except Exception as error:
    print("\033[91m", type(error).__name__, ": ", error, "\033[0m")

 RuntimeError :  Only Tensors of floating point and complex dtype can require gradients 

• When we create nn.Parameters, pytorch expects that parameter where it’s able to compute gradients on it.

• The issue is that with PyTorch, you can’t explicitly compute gradients on INT8 tensors.

• So above code snippet will give an error saying that only tensors of floating point and complex dtype can require gradients.

• So the right approach to save INT8 weights is instead of saving attributes as being an endless parameter, is to call the method called register buffer.

• This way instead of storing a parameter, we just store a buffer, meaning we don’t need to compute gradients on the tensor.

• You can initialize it with whatever dtype you want.

import torch
import torch.nn as nn
import torch.nn.functional as F

class W8A16LinearLayer(nn.Module):
    def __init__(self, in_features, out_features, 
                 bias=True, dtype=torch.float32):
        super().__init__()
        
        
        self.register_buffer(
            "int8_weights",
            torch.randint(
                -128, 127, (out_features, in_features), dtype=torch.int8
            )
        )
        
        self.register_buffer("scales", 
                             torch.randn((out_features), dtype=dtype)) # We are interested in inference only
        
        if bias:
            self.register_buffer("bias", 
                                 torch.randn((1, out_features), 
                                             dtype=dtype))
        
        else:
            self.bias = None
    def forward(self, input):
        return w8_a16_forward(self.int8_weights, 
                              input, self.scales, self.bias)

Quantize a Base Model

import torch
import torch.nn as nn
import torch.nn.functional as F

class W8A16LinearLayer(nn.Module):
    def __init__(self, in_features, out_features, 
                 bias=True, dtype=torch.float32):
        super().__init__()
        
        
        self.register_buffer(
            "int8_weights",
            torch.randint(
                -128, 127, (out_features, in_features), dtype=torch.int8
            )
        )
        
        self.register_buffer("scales", 
                             torch.randn((out_features), dtype=dtype))
        
        if bias:
            self.register_buffer("bias", 
                                 torch.randn((1, out_features), 
                                             dtype=dtype))
        
        else:
            self.bias = None

    def quantize(self, weights):
        w_fp32 = weights.clone().to(torch.float32)

        scales = w_fp32.abs().max(dim=-1).values / 127
        scales = scales.to(weights.dtype)

        int8_weights = torch.round(weights
                        /scales.unsqueeze(1)).to(torch.int8)

        self.int8_weights = int8_weights
        self.scales = scales
    
    def forward(self, input):
        return w8_a16_forward(self.int8_weights, 
                              input, self.scales, self.bias)

module = W8A16LinearLayer(4, 8)
print("Weights before:\n" , module.int8_weights)
random_matrix = torch.randn((4, 8), dtype=torch.bfloat16)
module.quantize(random_matrix)
print("Weights After:\n" , module.int8_weights)
print("Average quantization error:-",(random_matrix - module.int8_weights 
 * module.scales.unsqueeze(1)).abs().mean())

Weights before:
 tensor([[ -87,  -38,  -15,  -74],
        [ 107,  -95,    7,  -94],
        [-112, -110,  -21,  122],
        [ 112,  107,  -87,  105],
        [  67,  -67,   78,   78],
        [ -26,   -9,  -92,  -26],
        [ 126,   -3,   12,   17],
        [  51,  -30,  -28,   33]], dtype=torch.int8)
Weights After:
 tensor([[  87,  -75,   90,  -69,  -16,   52,  127,   21],
        [  52,   12,  -86,   40,   66,  -44,  -82,  127],
        [-108, -115,  -70,   92, -116,  127,  -10,   61],
        [  88,   71, -106,   56,  -98, -127, -106,  -78]], dtype=torch.int8)
Average quantization error:- tensor(0.0020, dtype=torch.bfloat16)

Replace PyTorch layers with Quantized Layers

• For language models, it better to not quantize the last layer.

import torch
import torch.nn as nn
import torch.nn.functional as F

def replace_linear_with_target(module, 
                               target_class, module_name_to_exclude):
    for name, child in module.named_children():
        if isinstance(child, nn.Linear) and not \
          any([x == name for x in module_name_to_exclude]):
            old_bias = child.bias

            new_module = target_class(child.in_features, 
                                      child.out_features, 
                                      old_bias is not None, 
                                      child.weight.dtype)
            setattr(module, name, new_module)
            if old_bias is not None:
              getattr(module, name).bias = old_bias
        else:
            # Recursively call the function for nested modules
            replace_linear_with_target(
                child, target_class, module_name_to_exclude)


def replace_linear_with_target_and_quantize(module, 
                               target_class, module_name_to_exclude):
    for name, child in module.named_children():
        if isinstance(child, nn.Linear) and not \
        any([x == name for x in module_name_to_exclude]):
            old_bias = child.bias
            old_weight = child.weight

            new_module = target_class(child.in_features, 
                                      child.out_features, 
                                      old_bias is not None, 
                                      child.weight.dtype)
            setattr(module, name, new_module) # current module is replaced by new_module

            getattr(module, name).quantize(old_weight)
            
            if old_bias is not None:
              getattr(module, name).bias = old_bias
        else:
            # Recursively call the function for nested modules
            replace_linear_with_target_and_quantize(child, 
                     target_class, module_name_to_exclude)

class DummyModel(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.emb = torch.nn.Embedding(1, 1)
    # Try with bias
    self.linear_1 = nn.Linear(1, 1)
    # Try without bias
    self.linear_2 = nn.Linear(1, 1, bias=False)
    # Lm prediction head
    self.lm_head = nn.Linear(1, 1, bias=False)

model_1 = DummyModel()
model_2 = DummyModel()
replace_linear_with_target_and_quantize(model_1, W8A16LinearLayer, ["lm_head"])
print("model_1",model_1)

replace_linear_with_target_and_quantize(model_2, W8A16LinearLayer, [])
print("model_2",model_2)

model_1 DummyModel(
  (emb): Embedding(1, 1)
  (linear_1): W8A16LinearLayer()
  (linear_2): W8A16LinearLayer()
  (lm_head): Linear(in_features=1, out_features=1, bias=False)
)
model_2 DummyModel(
  (emb): Embedding(1, 1)
  (linear_1): W8A16LinearLayer()
  (linear_2): W8A16LinearLayer()
  (lm_head): W8A16LinearLayer()
)

Quantize any Open Source PyTorch Model

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model_id = "Salesforce/codegen-350M-mono"

model = AutoModelForCausalLM.from_pretrained(model_id, 
                                    torch_dtype=torch.bfloat16, 
                                             low_cpu_mem_usage=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
print(pipe("def hello_world():", max_new_tokens=20, do_sample=False)[0]["generated_text"])
print("Model before:\n\n", model)
replace_linear_with_target_and_quantize(model, 
                                        W8A16LinearLayer, ["lm_head"])

print("Model after:\n\n", model)
print(pipe("def hello_world():", max_new_tokens=20, 
           do_sample=False)[0]["generated_text"])

Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.

def hello_world():
    print("Hello World")

hello_world()

# 파
Model before:

 CodeGenForCausalLM(
  (transformer): CodeGenModel(
    (wte): Embedding(51200, 1024)
    (drop): Dropout(p=0.0, inplace=False)
    (h): ModuleList(
      (0-19): 20 x CodeGenBlock(
        (ln_1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        (attn): CodeGenAttention(
          (attn_dropout): Dropout(p=0.0, inplace=False)
          (resid_dropout): Dropout(p=0.0, inplace=False)
          (qkv_proj): Linear(in_features=1024, out_features=3072, bias=False)
          (out_proj): Linear(in_features=1024, out_features=1024, bias=False)
        )
        (mlp): CodeGenMLP(
          (fc_in): Linear(in_features=1024, out_features=4096, bias=True)
          (fc_out): Linear(in_features=4096, out_features=1024, bias=True)
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.0, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=1024, out_features=51200, bias=True)
)

Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.

Model after:

 CodeGenForCausalLM(
  (transformer): CodeGenModel(
    (wte): Embedding(51200, 1024)
    (drop): Dropout(p=0.0, inplace=False)
    (h): ModuleList(
      (0-19): 20 x CodeGenBlock(
        (ln_1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        (attn): CodeGenAttention(
          (attn_dropout): Dropout(p=0.0, inplace=False)
          (resid_dropout): Dropout(p=0.0, inplace=False)
          (qkv_proj): W8A16LinearLayer()
          (out_proj): W8A16LinearLayer()
        )
        (mlp): CodeGenMLP(
          (fc_in): W8A16LinearLayer()
          (fc_out): W8A16LinearLayer()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.0, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=1024, out_features=51200, bias=True)
)
def hello_world():
    print("Hello World")

# hello_world()

# def hello_

• Above code snippet modifies the model inplace
• Also don’t try to change the lm_head otherwise it will not give the desired results
• All the rounding errors can sum up once you start generating a lot of tokens, until may be all of these errors get super large so that it affects the model’s performance

Load your Quantized Weights from HuggingFace Hub

• The idea is to quantize weights on bigger instance and then push it back to huggingface. So that we don’t have to load and quantize again and again.
• Then use meta device from pytorch to load the skeleton of the model instead of loading the whole model itself.
• Replace the original layers with the quantized layers
• Load the quantized weights from huggingfacehub

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "facebook/opt-125m"

model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)

replace_linear_with_target_and_quantize(model, 
                             W8A16LinearLayer, 
                                   ["lm_head"])

model

quantized_state_dict = model.state_dict()
torch.save(quantized_state_dict, r"/home/ritesh/github_repos/riteshrm.github.io/tmp/quantized_state_dict.pth")

How to upload on HF

from huggingface_hub import HfApi, create_repo

YOUR_HF_USERNAME = “” your_repo_id = f”{YOUR_HF_USERNAME}/opt-125m-quantized-dlai”

api = HfApi()

create_repo(your_repo_id)

api.upload_file( path_or_fileobj=“quantized_state_dict.pth”, path_in_repo=“quantized_state_dict.pth”, repo_id=your_repo_id )

import torch
from transformers import OPTForCausalLM, AutoTokenizer, AutoConfig

model_id = "facebook/opt-125m"
config = AutoConfig.from_pretrained(model_id)

with torch.device("meta"):
  model = OPTForCausalLM(config)

tokenizer = AutoTokenizer.from_pretrained(model_id)

for param in model.parameters():
  print(param)
  break

/home/ritesh/miniconda3/envs/quarto/lib/python3.10/site-packages/huggingface_hub/file_download.py:942: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(

Parameter containing:
tensor(..., device='meta', size=(50272, 768), requires_grad=True)

model

OPTForCausalLM(
  (model): OPTModel(
    (decoder): OPTDecoder(
      (embed_tokens): Embedding(50272, 768, padding_idx=1)
      (embed_positions): OPTLearnedPositionalEmbedding(2050, 768)
      (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (layers): ModuleList(
        (0-11): 12 x OPTDecoderLayer(
          (self_attn): OPTAttention(
            (k_proj): Linear(in_features=768, out_features=768, bias=True)
            (v_proj): Linear(in_features=768, out_features=768, bias=True)
            (q_proj): Linear(in_features=768, out_features=768, bias=True)
            (out_proj): Linear(in_features=768, out_features=768, bias=True)
          )
          (activation_fn): ReLU()
          (self_attn_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
          (fc1): Linear(in_features=768, out_features=3072, bias=True)
          (fc2): Linear(in_features=3072, out_features=768, bias=True)
          (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        )
      )
    )
  )
  (lm_head): Linear(in_features=768, out_features=50272, bias=False)
)

replace_linear_with_target(model, W8A16LinearLayer, ["lm_head"])
model

OPTForCausalLM(
  (model): OPTModel(
    (decoder): OPTDecoder(
      (embed_tokens): Embedding(50272, 768, padding_idx=1)
      (embed_positions): OPTLearnedPositionalEmbedding(2050, 768)
      (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (layers): ModuleList(
        (0-11): 12 x OPTDecoderLayer(
          (self_attn): OPTAttention(
            (k_proj): W8A16LinearLayer()
            (v_proj): W8A16LinearLayer()
            (q_proj): W8A16LinearLayer()
            (out_proj): W8A16LinearLayer()
          )
          (activation_fn): ReLU()
          (self_attn_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
          (fc1): W8A16LinearLayer()
          (fc2): W8A16LinearLayer()
          (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        )
      )
    )
  )
  (lm_head): Linear(in_features=768, out_features=50272, bias=False)
)

If loading from HF

from huggingface_hub import hf_hub_download

state_dict_cache_path = hf_hub_download( “ybelkada/opt-125m-quantized-dlai”, “quantized_state_dict.pth” )

state_dict = torch.load(r"/home/ritesh/github_repos/riteshrm.github.io/tmp/quantized_state_dict.pth")
model.load_state_dict(state_dict, strict=True, assign=True)

/tmp/ipykernel_1365/1695582919.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  state_dict = torch.load(r"/home/ritesh/github_repos/riteshrm.github.io/tmp/quantized_state_dict.pth")

<All keys matched successfully>

from transformers import pipeline

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
pipe("Hello today I am", max_new_tokens=40)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
pipe("Hello today I am giving a course about", max_new_tokens=10)

[{'generated_text': "Hello today I am giving a course about the history of the world's first ever solar eclipse"}]

Weights Packing

• Weights packing is important for storing quantized weights, because torch.int4 is not available as of today, so we need to store and load the weights in int8
• This is not ideal because:
  • tensor will occupy 8-bit per datapoint and might add a considerable overhead for large models
  • There would be no point of quantizing to 2/4 bits because we are still using 8-bit
• So, we need to pack values
• Consider a tensor with 4 values each with 2-bit(0,1,2,3) precision but stored in 8-bit
  • tensor = torch.tensor([1,0,3,2], dtype=torch.uint8)
  • 1:- 00000001
  • 0:- 00000000
  • 3:- 00000011
  • 2:- 00000010
• We can pack all these values into a single 8-bit value as 177
  • 177:- 10110001
• Advantages:-
  • It reflects the true memory footprint of the quantized weights Disadvantages:-
  • The unpacked tensors need to be a shape with a multiple of 8//nbits
  • It needs to unpack before performing an operation

Packing 2-bit Weights

import torch

# Example Tensor: [1, 0, 3, 2]
    # 1 0 3 2 - 01 00 11 10

    # Starting point of packed int8 Tensor
    # [0000 0000]

    ##### First Iteration Start:
    # packed int8 Tensor State: [0000 0000]
    # 1 = 0000 0001
    # 0000 0001
    # No left shifts in the First Iteration
    # After bit-wise OR operation between 0000 0000 and 0000 0001:
    # packed int8 Tensor State: 0000 0001
    ##### First Iteration End

    ##### Second Iteration Start:
    # packed int8 Tensor State: [0000 0001]
    # 0 = 0000 0000
    # 0000 0000
    # 2 left shifts:
    # [0000 0000] (1 shift)-> 0000 0000 (2 shift)-> 0000 0000
    # After bit-wise OR operation between 0000 0001 and 0000 0000:
    # packed int8 Tensor State: 0000 0001
    ##### Second Iteration End

    ##### Third Iteration Start:
    # packed int8 Tensor State: [0000 0001]
    # 3 = 0000 0011
    # 0000 0011
    # 4 left shifts:
    # [0000 0011] (1 shift)-> 0000 0110 (2 shift)-> 0000 1100
    # 0000 1100 (3 shift)-> 0001 1000 (4 shift)-> 0011 0000
    # After bit-wise OR operation between 0000 0001 and 0011 0000:
    # packed int8 Tensor State: 0011 0001
    ##### Third Iteration End

    ##### Fourth Iteration Start:
    # packed int8 Tensor State: [0011 0001]
    # 2 = 0000 0010
    # 0000 0010
    # 6 left shifts:
    # [0000 0010] (1 shift)-> 0000 0100 (2 shift)-> 0000 1000
    # 0000 1000 (3 shift)-> 0001 0000 (4 shift)-> 0010 0000
    # 0010 0000 (5 shift)-> 0100 0000 (6 shift)-> 1000 0000
    # After bit-wise OR operation between 0011 0001 and 1000 0000:
    # packed int8 Tensor State: 1011 0001
    ##### Fourth Iteration End

    # Final packed int8 Tensor State: [1011 0001]

def pack_weights(uint8tensor, bits):
    if uint8tensor.shape[0] * bits % 8 != 0:
        raise ValueError(f"The input shape needs to be a multiple \
        of {8 / bits} - got {uint8tensor.shape[0]}")

    num_values = uint8tensor.shape[0] * bits // 8

    num_steps = 8 // bits

    unpacked_idx = 0

    packed_tensor = torch.zeros((num_values), dtype=torch.uint8)

    # 1 0 3 2 - 01 00 11 10

    # [0000 0000] -> 0000 0001

    # 0000 0001

    # 0000 0000 - 0000 0000

    # 0000 0011 - 0011 0000 - 0011 0001

    # 1011 0001
    
    for i in range(num_values):
        for j in range(num_steps):
            packed_tensor[i] |= uint8tensor[unpacked_idx] << (bits * j)
            unpacked_idx += 1
    return packed_tensor
unpacked_tensor = torch.tensor([1, 0, 3, 2], 
                               dtype=torch.uint8)
pack_weights(unpacked_tensor, 2)
unpacked_tensor = torch.tensor([1, 0, 3, 2, 3, 3, 3, 3], 
                               dtype=torch.uint8)
pack_weights(unpacked_tensor, 2)