PyTorch Implementation about ringattention HOT 10 CLOSED

        num_q, _, _, _ = inputs.shape
        res = []
        for i in range(num_q):
            res.append(checkpoint(self.lm_head, inputs[i]))
        res = torch.stack(res, dim=0)

def lm_head(cell, hidden_states):
    outputs = cell(hidden_states)
    return outputs
outputs = checkpoint(lm_head, cell, hidden_states)

def scan(f, init, xs, length=None):
    if xs is None:
        xs = [None] * length
    carry = init
    ys = []
    for x in xs:
        carry, y = f(carry, x)
        ys.append(y)
    return carry, torch.stack(ys)

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint
from einops import rearrange

def scan(f, init, xs, length=None):
    if xs is None:
        xs = [None] * length
    carry = init
    ys = []
    for x in xs:
        carry, y = f(carry, x)
        ys.append(y)
    return carry, torch.stack(ys)

class FFN(nn.Module):
    def __init__(self, dim, hidden_dim, dropout):
        super().__init__()
        self.fc_in = nn.Linear(dim, hidden_dim)
        self.fc_out = nn.Linear(hidden_dim, dim)
        self.act = nn.GELU()
        self.resid_dropout = nn.Dropout(dropout)
        self.ln_2 = nn.LayerNorm(dim)

    def forward(self, hidden_states):
        hidden_states = self.ln_2(hidden_states)
        hidden_states = self.fc_in(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.fc_out(hidden_states)
        hidden_states = self.resid_dropout(hidden_states)
        return hidden_states

    def checkpoint_fn(self, carry, x):
        return None, checkpoint(self.forward, x)

def blockwise_compute_ffn(cell, inputs, chunk_size):
    inputs = rearrange(inputs, 'b (n c) d -> b n c d', c=chunk_size)
    inputs = rearrange(inputs, 'b n c d -> n b c d')
    _, res = scan(cell.checkpoint_fn, None, inputs, inputs.shape[0])
    res = rearrange(res, 'n b c d -> b (n c) d')
    return res

class Blockwise_LM_Head(nn.Module):
    def __init__(self, vocab_size, chunk_size):
        super().__init__()
        self.vocab_size = vocab_size
        self.chunk_size = chunk_size
        self.lm_head = nn.Linear(
            chunk_size, 
            vocab_size, 
            bias=True
        )

    def checkpoint_fn(self, carry, x):
        return None, checkpoint(self.lm_head, x)

    def forward(self, inputs):
        inputs = rearrange(inputs, 'b (n c) d -> b n c d', c=self.chunk_size)
        inputs = rearrange(inputs, 'b n c d -> n b c d')
        _, res = scan(self.checkpoint_fn, None, inputs, inputs.shape[0])
        res = rearrange(res, 'n b c d -> b (n c) d')
        return res

def cross_entropy_loss_and_accuracy(logits, tokens, valid):
    valid_text_length = torch.clamp(torch.sum(valid, axis=-1), min=1e-10)

    token_log_prob = torch.log_softmax(logits, axis=-1)
    token_log_prob = torch.gather(token_log_prob, -1, tokens.unsqueeze(-1)).squeeze(-1)

    token_log_prob = torch.where(valid > 0.0, token_log_prob, torch.tensor(0.0, device=logits.device))
    correct = torch.where(
        valid > 0.0,
        logits.argmax(axis=-1) == tokens,
        torch.tensor(False, device=logits.device)
    )
    return token_log_prob, correct, valid_text_length

def _loss_and_accuracy(carry, args):
    loss, accuracy, num = carry
    logits, tokens, valid = args
    token_log_prob, correct, valid_text_length = checkpoint(cross_entropy_loss_and_accuracy, logits, tokens, valid)
    loss = loss + torch.sum(token_log_prob, axis=-1) / valid_text_length
    accuracy = accuracy + torch.sum(correct.float(), axis=-1) / valid_text_length
    num = num + 1
    return (loss, accuracy, num), None

def blockwise_cross_entropy(logits, tokens, valid=None, chunk_size=None):
    if valid is None:
        valid = torch.ones(tokens.shape[:2])
    valid = valid.float()
    logits = logits.view(-1, logits.shape[-1])
    tokens = tokens.view(-1,)
    valid = valid.view(-1,)

    num_chunk = logits.shape[0] // chunk_size
    logits = rearrange(logits, '(n c) d -> n c d', c=chunk_size)
    tokens = rearrange(tokens, '(n c) -> n c', c=chunk_size)
    valid = rearrange(valid, '(n c) -> n c', c=chunk_size)
    (loss, accuracy, num), _ = scan(
        _loss_and_accuracy, (torch.tensor(0.0, device=logits.device), torch.tensor(0.0, device=logits.device), 0), 
        (logits, tokens, valid), length=num_chunk
    )
    loss = - loss / num
    return loss, accuracy