Scaling Vision Transformers to 22 billion parameters

Mohit_Kumar: <pre><code class="lang-auto">import triton
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.profiler import record_function
from typing import Type



class LayerScale(nn.Module):
  def __init__(self, dim: int, init_values: float = 1e-5, inplace = False):
    super().__init__()
    self.gamma = nn.Parameter(torch.ones(dim) * init_values)
    self.inplace = inplace

  def forward(self, x: torch.Tensor):
    return x.mul_(self.gamma) if self.inplace else x * self.gamma


class ParallelBlockMultiHeadAttention(nn.Module):
    def __init__(
            self,
            dim,
            num_heads,
            norm_layer: Type[nn.Module] = nn.LayerNorm
            ):
        super().__init__()
        assert dim % num_heads == 0, "dim should be multiple of num_heads"
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads

        self.q_norm = norm_layer(self.head_dim)
        self.k_norm = norm_layer(self.head_dim)

    def forward(self, QKV):
        B, N, C = QKV.shape
        assert C // 3 == self.dim, f"Inner dim of QKV {C} != dim = {self.dim}"
        # B, N, 3, num_heads, head_dim =&gt; 3, B, num_heads, N, head_dim
        # Is it bad to use view here?
        QKV = QKV.view(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = QKV.unbind(0)

        # These excessively large logits, after being passed through a softmax function, 
        # led to near-zero entropy attention weights that were "almost one-hot"
        # Prior to QK Normalisation, it was often necessary to reduce the learning rate for larger ViT 
        # models (e.g., from 1e-3 down to 4e-4 for ViT-H) to maintain stability.
        # With QK Normalisation, ViT-22B can maintain a higher and consistent learning rate of 1e-3
        q, k = self.q_norm(q), self.k_norm(k)

        attn = torch.matmul(q, k.permute(0, 1, 3, 2))
        attn = attn * self.head_dim ** -0.5
        attn = attn.softmax(dim=-1)
        x = attn @ v
        return x


class ParallelBlock(nn.Module):
    def __init__(
            self, 
            dim, 
            num_heads,
            mlp_ratio: float = 0.7,
            act_layer: Type[nn.Module] = nn.GELU, # Will discuss what nn.GELU is
            ):
        super().__init__()

        self.dim = dim
        self.num_heads = num_heads
        assert dim % num_heads == 0, "dim should be multiple of num_heads"
        self.head_dim = dim // num_heads

        self.mlp_hidden_dim = int(dim * mlp_ratio)

        self.norm = nn.LayerNorm(dim)
        out_fusion1 = dim * 3 + self.mlp_hidden_dim
        self.fused_proj1 = nn.Linear(in_features = dim, out_features = out_fusion1)

        self.attn = ParallelBlockMultiHeadAttention(dim=dim, num_heads=num_heads)
        self.attn_out_proj = nn.Linear(in_features = dim, out_features = dim)

        self.mlp_activation = act_layer()
        self.mlp_out_proj = nn.Linear(in_features = self.mlp_hidden_dim, out_features = self.dim)

        self.fused_proj2 = nn.Linear(in_features = 2 * dim, out_features = dim)
        self.layer_scale = LayerScale(dim)
 
    def forward(self, x):
        _, _, C = x.shape
        assert C == self.dim, "inner dimension of x should be equal to the \
           out_features of fused_linear_layer1"

        norm_x = self.norm(x)

        fused_proj1 = self.fused_proj1(norm_x)
        
        qkv, mlp_hidden = torch.split(
            fused_proj1, 
            [3*self.dim, self.mlp_hidden_dim], 
            dim=-1
        )
        
        stream1 = torch.cuda.Stream()
        stream2 = torch.cuda.Stream()
        current_stream = torch.cuda.current_stream()
        with torch.cuda.stream(stream1):
            print("start1")
            attn_out = self.attn(qkv).permute(0, 2, 1, 3).flatten(-2)
            attn_out = self.attn_out_proj(attn_out)
            print("end1")
        with torch.cuda.stream(stream2):
            print("start2")
            mlp_out = self.mlp_activation(mlp_hidden)
            mlp_out = self.mlp_out_proj(mlp_out)
            print("end2")

        current_stream.wait_stream(stream1)
        current_stream.wait_stream(stream2)
        torch.cuda.synchronize()

        combined = torch.cat([attn_out, mlp_out], dim=-1)
        fused_out = self.fused_proj2(combined)

        return x + self.layer_scale(fused_out)


if __name__ == "__main__":
    B, N, C = 4, 1024, 4096
    num_heads = 4
    x = torch.randn((B, N, C))
    block = ParallelBlock(dim=C, num_heads=num_heads)
    output = block(x)
    assert output.shape == x.shape
</code></pre>
<p>The code I have provided is model architecture of <a href="https://arxiv.org/abs/2302.05442" rel="noopener nofollow ugc">Scaling Vision Transformers to 22 Billion Parameters</a>, the paper uses JAX and FLAX. I am new to parallel programming. As you can see in the code, I tried to launch 2 modules independently, but <code>cuda.Stream</code> will queue these streams and there will be no parallel execution of these two modules.<br>
Output:</p>
<pre><code class="lang-auto">start1
end1
start2
end2
</code></pre>
<p>Q1. Is there any way to launch these modules in parallel then use cuda sync to wait (that’s other logic…), even if I wraps the streams in another module?<br>
Q2. If not then I have to go for multiple devices. But I am curious how they achieved this (launching two modules in parallel) in paper using JAX (may be) but we can’t do same in pytorch (not so sure)?</p>

ptrblck: <p><a href="https://discuss.pytorch.org/t/running-independent-nn-module-instances-in-nn-modulelist-truly-in-parallel-in-pytorch/188166/3">This post</a> as well as <a href="https://discuss.pytorch.org/t/running-two-batches-in-parallel-using-cuda-streams-does-not-overlap-during-training/215369/2">this one</a> might be helpful.</p>