Class Stack

transformer.py:323–381

    
@si_module
class Stack(nn.Module):
    class Config:
        layers: int
        dim: int
        seq_len: int
        n_head: int = 32
        ff_dim: int = None
        kv_heads: int = None
        eps: float = 1e-5
        theta: Union[int, float] = 10_000
        causal: bool = True

        from_pretrained: Optional[Tuple[str, int]] = None

    def __init__(self, c: Config):
        super().__init__()

        from_pretrained = c.from_pretrained
        if exists(from_pretrained):
            checkpoint = load_ckpt(c.from_pretrained)

        self.shape_rotator = ShapeRotator(c.dim//c.n_head, c.seq_len, theta=c.theta)

        self.layers = nn.ModuleList([
            Block(
                dim=c.dim,
                layer_id=l,
                n_head=c.n_head,
                kv_heads=c.kv_heads,
                ff_dim=c.ff_dim,
                eps=c.eps,
                causal=c.causal,
                shape_rotator=self.shape_rotator,
            ) for l in range(c.layers)
        ])

        kv_heads = c.kv_heads or c.n_head
        head_dim = c.dim // c.n_head
        cache_shape = [c.layers, c.seq_len, 2, kv_heads, head_dim]
        self.cache_shape = cache_shape
        self.cache = [None] * c.layers

        if exists(from_pretrained):
            self.load_state_dict(checkpoint)

    def init_cache(self, bsize, device, dtype, length:int=None):
        if self.cache_shape is None:
            return
        cache_shape = self.cache_shape.copy()
        cache_shape[1] = length or cache_shape[1]
        self.cache = T.full((bsize, *cache_shape), CACHE_FILL_VALUE, device=device, dtype=dtype).transpose(0, 1)

    def deinit_cache(self):
        self.cache = [None] * len(self.cache)

    def forward(self, x: Tensor) -> Tensor:
        for l, layer in enumerate(self.layers):
            x = layer(x, kv=self.cache[l])