<jupyter_start><jupyter_text><jupyter_code>!nvidia-smi
import torch
import torch.nn as nn
import torch.autograd.profiler as profiler<jupyter_output><empty_output><jupyter_text>Both modules are stripped versions of lucidrains code<jupyter_code>class FeedForward(nn.Module):
    def __init__(self, d_model, p=0.):
        super().__init__()
        self.lin1 = nn.Linear(d_model, d_model*4)
        self.act = nn.GELU()
        self.drop = nn.Dropout(p)
        self.lin2 = nn.Linear(d_model*4, d_model)

    def forward(self, x):
        x = self.lin1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.lin2(x)
        return x
class Chunk(nn.Module):
    def __init__(self, chunks, func, dim=1):
        super().__init__()
        self.dim = dim
        self.chunks = chunks
        self.func = func

    def forward(self, x):
        if self.chunks == 1:
            return self.fn(x)
        chunks = x.chunk(self.chunks, dim = self.dim)
        return torch.cat([self.func(c) for c in chunks], dim = self.dim)<jupyter_output><empty_output><jupyter_text>sanity check<jupyter_code># config
bs = 1
sl = 2**8
d_model = 768
p = 0.
x = torch.randn(bs, sl, d_model)
ff = FeedForward(d_model, p)
chunk_ff = Chunk(10, ff)
with torch.no_grad():
    out = ff(x)
    c_out = chunk_ff(x)
assert ((out - c_out).abs() < 1e-5).all()
del x, ff, chunk_ff, out, c_out<jupyter_output><empty_output><jupyter_text>test 1 forward pass + no_grad  vanilla ff<jupyter_code>bs = 8
sl = 2**16
d_model = 1024
p = 0.
x = torch.randn(bs, sl, d_model, device='cuda')
!nvidia-smi
ff = FeedForward(d_model, p).cuda()
with profiler.profile(record_shapes=True, profile_memory=True, use_cuda=True) as prof:
    with torch.no_grad():
        out = ff(x)
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))<jupyter_output><empty_output><jupyter_text>chunked ff<jupyter_code>bs = 8
sl = 2**16
d_model = 1024
p = 0.
x = torch.randn(bs, sl, d_model, device='cuda')
!nvidia-smi
ff = FeedForward(d_model, p).cuda()
chunk_ff = Chunk(10, ff)
with profiler.profile(record_shapes=True, profile_memory=True, use_cuda=True) as prof:
    with torch.no_grad():
        out = chunk_ff(x)
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))<jupyter_output>----------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                  Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  
----------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
           aten::empty        38.83%       4.553ms        38.83%       4.553ms     111.042us       0.000us         0.00%       0.000us       0.000us           0 b           0 b      20.02 Gb      20.02 Gb            41  
          aten::matmul         3.44%     403.090us        41.27%       4.838ms     241.905us     188.290us        [...]<jupyter_text>test 2 forward pass  vanilla ff<jupyter_code># config
bs = 6
sl = 2**16
d_model = 1024
p = 0.
x = torch.randn(bs, sl, d_model).cuda()
!nvidia-smi
ff = FeedForward(d_model, p).cuda()
with profiler.profile(record_shapes=True, profile_memory=True, use_cuda=True) as prof:
    out = ff(x)
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))<jupyter_output><empty_output><jupyter_text>chunked ff<jupyter_code># config
bs = 6
sl = 2**16
d_model = 1024
p = 0.
x = torch.randn(bs, sl, d_model).cuda()
!nvidia-smi
ff = FeedForward(d_model, p).cuda()
chunk_ff = Chunk(10, ff)
with profiler.profile(record_shapes=True, profile_memory=True, use_cuda=True) as prof:
    out = chunk_ff(x)
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))<jupyter_output><empty_output><jupyter_text>test fwd + bwd  ff<jupyter_code># config
bs = 4
sl = 2**16
d_model = 1024
p = 0.
x = torch.randn(bs, sl, d_model).cuda()
!nvidia-smi
ff = FeedForward(d_model, p).cuda()
with profiler.profile(record_shapes=True, profile_memory=True, use_cuda=True) as prof:
    out = ff(x)
    out.sum().backward()
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))<jupyter_output><empty_output><jupyter_text>chunked ff<jupyter_code># config
bs = 4
sl = 2**16
d_model = 1024
p = 0.
x = torch.randn(bs, sl, d_model).cuda()
!nvidia-smi
ff = FeedForward(d_model, p).cuda()
chunk_ff = Chunk(10, ff)
with profiler.profile(record_shapes=True, profile_memory=True, use_cuda=True) as prof:
    out = chunk_ff(x)
    out.sum().backward()
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.2006465952717721, 'fraction_numerical': 0.03475449585774904, 'mean_word_length': 2.42797783933518, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '421642', 'license': 'MIT', 'n_tokens_mistral': 1980, 'n_tokens_neox': 1829, 'n_words': 369, 'path': 'exploration/chunked_feed_forward_mem_test.ipynb', 'repo_name': 'tyoc213-contrib/reformer-fastai-old', 'size': '4949'}