How to benchmark training rig with Nvidia RTX3080Ti

Palkos: <p>Hi</p>
<p>I am currently trying to train a very simple transformer based on this codebase with some minor changes:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/karpathy/ng-video-lecture">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/karpathy/ng-video-lecture" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/344;"><img src="https://opengraph.githubassets.com/7637730b26794259c1ca3ac66866540c6df77866f392742522df2663844d734c/karpathy/ng-video-lecture" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/karpathy/ng-video-lecture" target="_blank" rel="noopener nofollow ugc">GitHub - karpathy/ng-video-lecture</a></h3>

  <p>Contribute to karpathy/ng-video-lecture development by creating an account on GitHub.</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>The forward pass for my data seems to be very quick (0.01s) but backwards pass is externally slow (about 7s on average). This is very concerning, I expect the time to be 2, maybe 3 times the forward pass but not so long.</p>
<p>I had to change some parameters to fit the requirements:<br>
Vocabulary size is 256<br>
The sequence length is 1024<br>
Batch size is set to 64<br>
Heads and blocks are both to 4 to fit this to my GPU.<br>
Feed forward layer size is reduce to 256</p>
<p>I am running this on Intel i7-7700K, 16GB of ram and Nvidia 3080Ti with on Windows.<br>
Python is 3.10.7, pytorch 2.3,  cuda 12.1</p>
<p>So perhaps the performance is fine and it is simply going to be that slow. But is there a way to check that the setup is working optimally? Can this train faster or is that pretty much 100% and I just should be patient?</p>
<p>Here is the code I run with random data to check the performance:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import numpy as np
from torch.nn import functional as F
from datetime import datetime
from tqdm import trange

# hyperparameters
batch_size = 64 # how many independent sequences will we process in parallel?
block_size = 1024 # what is the maximum context length for predictions?
max_iters = 500000
eval_interval = 1
learning_rate = 3e-4

eval_iters = 200
n_embd = 256
n_head = 4
n_layer = 4
dropout = 0.2

vocab_size = 256
# ------------

# data loading
def get_batch(split):
    # generate a small batch of data of inputs x and targets y
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out


class Tokenizer():
    
    binRanges = []
    binValues = []

    def __init__(self, priceDiffRange, binsCount, scalingFactor):
        oneSide = (int)(binsCount/2);
        
        for iter in range(oneSide):
            self.binRanges.append(priceDiffRange)
            self.binRanges.append(-priceDiffRange)
            priceDiffRange = priceDiffRange/ scalingFactor

        self.binRanges.sort()
        
        self.binValues.append(-priceDiffRange)

        for iter in range(1, binsCount-1):
            self.binValues.append((self.binRanges[iter] + self.binRanges[iter-1])/2)
        
        self.binValues.append(priceDiffRange)

    def tokenize(self, raw_data):
        
        result = []

        for val in raw_data:
            token = -1;
            for iter in range(vocab_size-1):
                if(self.binRanges[iter]&gt; val):
                    token = iter;
                    break
            
            if token ==-1:
                token = vocab_size-1
            
            result.append(token)
            
        return result


    def detokenize(self, tokens):
          
        result = []

        for val in tokens:
            if(val&gt;vocab_size-1) : 
                result.append(self.binValues[vocab_size-1])
            else:
                result.append(self.binValues[val])
            
        return result

class Head(nn.Module):
    """ one head of self-attention """

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # input of size (batch, time-step, channels)
        # output of size (batch, time-step, head size)
        B,T,C = x.shape
        k = self.key(x)   # (B,T,hs)
        q = self.query(x) # (B,T,hs)
        # compute attention scores ("affinities")
        wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5 # (B, T, hs) @ (B, hs, T) -&gt; (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
        wei = F.softmax(wei, dim=-1) # (B, T, T)
        wei = self.dropout(wei)
        # perform the weighted aggregation of the values
        v = self.value(x) # (B,T,hs)
        out = wei @ v # (B, T, T) @ (B, T, hs) -&gt; (B, T, hs)
        return out

class MultiHeadAttention(nn.Module):
    """ multiple heads of self-attention in parallel """

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(head_size * num_heads, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

class FeedFoward(nn.Module):
    """ a simple linear layer followed by a non-linearity """

    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

class Block(nn.Module):
    """ Transformer block: communication followed by computation """

    def __init__(self, n_embd, n_head):
        # n_embd: embedding dimension, n_head: the number of heads we'd like
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedFoward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

class GPTLanguageModel(nn.Module):

    def __init__(self):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd) # final layer norm
        self.lm_head = nn.Linear(n_embd, vocab_size)

        # better init, not covered in the original GPT video, but important, will cover in followup video
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        # idx and targets are both (B,T) tensor of integers
        tok_emb = self.token_embedding_table(idx) # (B,T,C)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
        x = tok_emb + pos_emb # (B,T,C)
        x = self.blocks(x) # (B,T,C)
        x = self.ln_f(x) # (B,T,C)
        logits = self.lm_head(x) # (B,T,vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        # idx is (B, T) array of indices in the current context
        for _ in range(max_new_tokens):
            # crop idx to the last block_size tokens
            idx_cond = idx[:, -block_size:]
            # get the predictions
            logits, loss = self(idx_cond)
            # focus only on the last time step
            logits = logits[:, -1, :] # becomes (B, C)
            # apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1) # (B, C)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            # append sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
        return idx
    


################################################################################
########  Main
################################################################################
    
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)

torch.manual_seed(1983)

# random sample datasets (for test if we do not want to wait for tokenizer)
val_data = torch.tensor(np.random.rand(100000), dtype=torch.long)
train_data = torch.tensor(np.random.rand(100000), dtype=torch.long)


model = GPTLanguageModel()
m = model.to(device)
# print the number of parameters in the model
print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')

# create a PyTorch optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

lastLoss = float("inf")
with trange(max_iters, desc="Training Progress", unit="batches") as pbar:
    for iter in pbar:

        # every once in a while evaluate the loss on train and val sets
        # if iter % eval_interval == 0 or iter == max_iters - 1:
        #     losses = estimate_loss()		
        #    print(datetime.now().strftime("%H:%M:%S"))
             #   print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
		
        # context = torch.zeros((1, 1), dtype=torch.long, device=device)
        # print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))

        # sample a batch of data
        xb, yb = get_batch('train')

        # evaluate the loss
        logits, loss = model(xb, yb)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()        
        optimizer.step()
        
        pbar.set_postfix(loss=loss.item())

# generate from the model
# context = torch.zeros((1, 1), dtype=torch.long, device=device)
# print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))
#open('more.txt', 'w').write(decode(m.generate(context, max_new_tokens=10000)[0].tolist()))

</code></pre>

ptrblck: <p>I don’t see how you are profiling your forward and backward passes, but note that CUDA operations are executed asynchronously w.r.t the CPU. If you are using host timers you would thus need to synchronize the code before starting and stopping the timers. Otherwise your profiles would be misleading as synchronizing operations, such as <code>.item()</code> calls, will accumulate the time of previous operations.</p>

Palkos: <p>I have used simple host based timers to be fair (which are not included in the sample code), however the problem still remains. The back prop took way more time than model. Are you suggesting that the return from model call is async and then back prop one has to wait for the model to finish before it could be calculated?</p>
<p>Could you suggest some good method to benchmark it?</p>

ptrblck: <aside class="quote no-group" data-username="Palkos" data-post="3" data-topic="202132">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/958977/48.png" class="avatar"> Palkos:</div>
<blockquote>
<p>Could you suggest some good method to benchmark it?</p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="202132">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>If you are using host timers you would thus need to synchronize the code before starting and stopping the timers.</p>
</blockquote>
</aside>
<p>Synchronize the code via <code>torch.cuda.synchronize()</code> or use the <code>torch.utils.benchmark</code> utils, which will add warmup iterations and will synchronize the code for you.<br>
You can also use a visual profiler, such as Nsight Systems, to understand the execution time in more detail.</p>

Palkos: <p>Many thanks for help.<br>
I tried to use the torch.utils.benchmark and I got some interesting results.</p>
<p>Firstly I was not able to run with exactly this same setup as originally (sequence length 1024 and batch size of 64) as cuda was not able to allocate the GPU memory. Seems like adding benchmark pushed the GPU memory over the limit. Then I tried to reduce the memory usage with set of experiments:</p>
<p>batch size=64<br>
sequence size = 128</p>
<p>Forward Test<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x0000027FFA99B0D0&gt;<br>
forward_test(xb, yb)<br>
setup: from <strong>main</strong> import forward_test<br>
9.59 ms<br>
1 measurement, 100 runs , 1 thread<br>
BackwardTest<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000001F14853BB80&gt;<br>
backward_test(loss)<br>
setup: from <strong>main</strong> import backward_test<br>
13.16 ms<br>
1 measurement, 100 runs , 1 thread</p>
<hr>
<p>batch size=64<br>
sequence size = 256</p>
<p>Forward Test<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x0000027FFA99B0D0&gt;<br>
forward_test(xb, yb)<br>
setup: from <strong>main</strong> import forward_test<br>
15.44 ms<br>
1 measurement, 100 runs , 1 thread<br>
BackwardTest<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x0000027FFA99B0D0&gt;<br>
backward_test(loss)<br>
setup: from <strong>main</strong> import backward_test<br>
28.24 ms<br>
1 measurement, 100 runs , 1 thread</p>
<hr>
<p>batch size=64<br>
sequence size = 512</p>
<p>Forward Test<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000002A76E5EB0A0&gt;<br>
forward_test(xb, yb)<br>
setup: from <strong>main</strong> import forward_test<br>
38.34 ms<br>
1 measurement, 100 runs , 1 thread<br>
BackwardTest<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000002A76E5EB0A0&gt;<br>
backward_test(loss)<br>
setup: from <strong>main</strong> import backward_test<br>
69.13 ms<br>
1 measurement, 100 runs , 1 thread</p>
<hr>
<p>batch size=64<br>
sequence size = 768</p>
<p><strong>Forward Test</strong><br>
<strong>&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000002772513B0A0&gt;</strong><br>
<strong>forward_test(xb, yb)</strong><br>
<strong>setup: from <strong>main</strong> import forward_test</strong><br>
<strong>5.79 s</strong><br>
**  1 measurement, 100 runs , 1 thread**<br>
<strong>BackwardTest</strong><br>
<strong>&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000002772513B0A0&gt;</strong><br>
<strong>backward_test(loss)</strong><br>
<strong>setup: from <strong>main</strong> import backward_test</strong><br>
**  124.88 ms**<br>
**  1 measurement, 100 runs , 1 thread**</p>
<p><strong>Forward Test</strong><br>
<strong>&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000002978FC3B0D0&gt;</strong><br>
<strong>forward_test(xb, yb)</strong><br>
<strong>setup: from <strong>main</strong> import forward_test</strong><br>
**  6.50 s**<br>
**  1 measurement, 100 runs , 1 thread**<br>
<strong>BackwardTest</strong><br>
<strong>&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000002978FC3B0D0&gt;</strong><br>
<strong>backward_test(loss)</strong><br>
<strong>setup: from <strong>main</strong> import backward_test</strong><br>
**  136.64 ms**<br>
**  1 measurement, 100 runs , 1 thread**</p>
<hr>
<p>batch size = 4<br>
sequence = 1024</p>
<p>Forward Test<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000001D01403B0A0&gt;<br>
forward_test(xb, yb)<br>
setup: from <strong>main</strong> import forward_test<br>
9.22 ms<br>
1 measurement, 100 runs , 1 thread<br>
BackwardTest<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000001D01403B0A0&gt;<br>
backward_test(loss)<br>
setup: from <strong>main</strong> import backward_test<br>
17.46 ms<br>
1 measurement, 100 runs , 1 thread</p>
<hr>
<p>batch size = 16<br>
sequence = 1024</p>
<p>Forward Test<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000001D01403B0A0&gt;<br>
forward_test(xb, yb)<br>
setup: from <strong>main</strong> import forward_test<br>
9.22 ms<br>
1 measurement, 100 runs , 1 thread<br>
BackwardTest<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000001D01403B0A0&gt;<br>
backward_test(loss)<br>
setup: from <strong>main</strong> import backward_test<br>
17.46 ms<br>
1 measurement, 100 runs , 1 thread</p>
<hr>
<p>batch size = 8<br>
sequence = 2048</p>
<p>Forward Test<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000001930BE3B0D0&gt;<br>
forward_test(xb, yb)<br>
setup: from <strong>main</strong> import forward_test<br>
52.10 ms<br>
1 measurement, 100 runs , 1 thread<br>
BackwardTest<br>
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x000001930BE3B0D0&gt;<br>
backward_test(loss)<br>
setup: from <strong>main</strong> import backward_test<br>
85.83 ms<br>
1 measurement, 100 runs , 1 thread</p>
<p>The interesting thing is that I can go to 8 batch  and 2048 sequence length (which is the one I wanted to achieve originally) and run it very fast.</p>
<p>However - if you look at the sequence length of 768 (marked in bold above) - the forward run is getting crazy. Nearly 6s!<br>
I repeated that test few times and results are the same. Any suggestion why that is the case?</p>