why there is massive memory overhead in pytorhc objects , for example in this code the saved model is onlt 3.5gb on disk but on gpu much more , also each btch takes a lot of memory evenn though it's just text : import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer
from datasets import load_dataset
from tqdm import tqdm
import time
import os
import math

# Constants
BATCH_SIZE = 16
SEQUENCE_LENGTH = 64
LEARNING_RATE = 2e-5
NUM_EPOCHS = 3
SAVE_STEPS = 2000
EMBEDDING_DIM = 4096
NUM_LAYERS = 16
HEADS = 16
MEMORY_SIZE = 256
MODEL_PATH = "reversible_transformer_lm.pth"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load GPT-2 tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

class ReversibleBlock(nn.Module):
    def __init__(self, f_block, g_block):
        super().__init__()
        self.f = f_block
        self.g = g_block

    def forward(self, x1, x2):
        y1 = x1 + self.f(x2)
        y2 = x2 + self.g(y1)
        return y1, y2

class ExternalMemory(nn.Module):
    def __init__(self, memory_size, embed_dim):
        super(ExternalMemory, self).__init__()
        self.memory_size = memory_size
        self.embed_dim = embed_dim
        self.memory = nn.Parameter(torch.randn(memory_size, embed_dim) / math.sqrt(embed_dim))
        self.controller = nn.Linear(embed_dim * 2, embed_dim)

    def read(self, query):
        scores = torch.matmul(query, self.memory.t())
        attention_weights = F.softmax(scores, dim=-1)
        read_content = torch.matmul(attention_weights, self.memory)
        return read_content, attention_weights

    def write(self, key, value):
        scores = torch.matmul(key, self.memory.t())
        attention_weights = F.softmax(scores, dim=-1)
        delta = torch.matmul(attention_weights.unsqueeze(-1), value.unsqueeze(1))
        self.memory.data = self.memory.data + delta.mean(dim=0)
        self.memory.data = F.normalize(self.memory.data, p=2, dim=-1)

    def forward(self, x):
        B, T, E = x.size()
        outputs = []
        for t in range(T):
            query = x[:, t, :]
            read_content, _ = self.read(query)
            combined = torch.tanh(self.controller(torch.cat([query, read_content], dim=-1)))
            outputs.append(combined.unsqueeze(1))
            self.write(query, combined)
        outputs = torch.cat(outputs, dim=1)
        return outputs

class DynamicAttention(nn.Module):
    def __init__(self, attention_module, embed_dim):
        super(DynamicAttention, self).__init__()
        self.attention_module = attention_module
        self.gate = nn.Linear(embed_dim, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        attn_output = self.attention_module(x)
        gate_values = self.sigmoid(self.gate(x))
        output = attn_output * gate_values
        return output

class CausalLinearAttention(nn.Module):
    def __init__(self, embed_dim, heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.heads = heads
        self.head_dim = embed_dim // heads
        assert self.head_dim * heads == embed_dim, f"embed_dim {embed_dim} must be divisible by heads {heads}"

        self.to_q = nn.Linear(embed_dim, embed_dim)
        self.to_k = nn.Linear(embed_dim, embed_dim)
        self.to_v = nn.Linear(embed_dim, embed_dim)
        self.to_out = nn.Linear(embed_dim, embed_dim)
        self.norm = nn.LayerNorm(embed_dim)
        self.init_weights()

    def init_weights(self):
        nn.init.xavier_uniform_(self.to_q.weight, gain=1/math.sqrt(2))
        nn.init.xavier_uniform_(self.to_k.weight, gain=1/math.sqrt(2))
        nn.init.xavier_uniform_(self.to_v.weight, gain=1/math.sqrt(2))
        nn.init.xavier_uniform_(self.to_out.weight, gain=1/math.sqrt(2))
        nn.init.zeros_(self.to_q.bias)
        nn.init.zeros_(self.to_k.bias)
        nn.init.zeros_(self.to_v.bias)
        nn.init.zeros_(self.to_out.bias)

    def forward(self, x):
        B, T, E = x.size()
        x = self.norm(x)

        Q = self.to_q(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)
        K = self.to_k(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)
        V = self.to_v(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)

        causal_mask = torch.tril(torch.ones(T, T, device=x.device)).unsqueeze(0).unsqueeze(0)
        scores = scores.masked_fill(causal_mask == 0, float('-inf'))

        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)

        out = out.transpose(1, 2).contiguous().view(B, T, E)
        return self.to_out(out)

class FeedForward(nn.Module):
    def __init__(self, embed_dim, expansion_factor=4):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * expansion_factor),
            nn.GELU(),
            nn.Linear(embed_dim * expansion_factor, embed_dim)
        )
        self.norm = nn.LayerNorm(embed_dim)
        self.init_weights()

    def init_weights(self):
        for layer in self.net:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_uniform_(layer.weight, gain=1/math.sqrt(2))
                nn.init.zeros_(layer.bias)

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

class ReversibleTransformerLayer(nn.Module):
    def __init__(self, embed_dim, heads):
        super().__init__()
        embed_dim_half = embed_dim // 2
        self.attn_block = DynamicAttention(CausalLinearAttention(embed_dim_half, heads), embed_dim_half)
        self.ffn_block = FeedForward(embed_dim_half)
        self.rev_block = ReversibleBlock(self.attn_block, self.ffn_block)

    def forward(self, x1, x2):
        return self.rev_block(x1, x2)

class ConfidencePredictor(nn.Module):
    def __init__(self, embed_dim, vocab_size):
        super().__init__()
        self.predictor = nn.Linear(embed_dim, 1)
        self.loss_fn = nn.BCEWithLogitsLoss()

    def forward(self, hidden_states, labels):
        logits = self.predictor(hidden_states).squeeze(-1)
        with torch.no_grad():
            predictions = hidden_states.argmax(dim=-1)
            correct = (predictions == labels).float()
        loss = self.loss_fn(logits, correct)
        confidence_scores = torch.sigmoid(logits)
        return confidence_scores, loss

class ReversibleTransformer(nn.Module):
    def __init__(self, depth, embed_dim, heads, vocab_size, memory_size):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.pos_embed = nn.Parameter(torch.zeros(1, SEQUENCE_LENGTH + 1, embed_dim))
        self.global_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.memory_module = ExternalMemory(memory_size, embed_dim)
        self.depth = depth
        self.layers = nn.ModuleList([
            ReversibleTransformerLayer(embed_dim, heads)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)
        self.out = nn.Linear(embed_dim, vocab_size)
        self.confidence_predictor = ConfidencePredictor(embed_dim, vocab_size)
        self.init_weights()

    def init_weights(self):
        nn.init.normal_(self.embed.weight, mean=0, std=0.02)
        nn.init.normal_(self.pos_embed, mean=0, std=0.01)
        nn.init.normal_(self.global_token, mean=0, std=0.01)
        nn.init.xavier_uniform_(self.out.weight, gain=1/math.sqrt(2))
        nn.init.zeros_(self.out.bias)

    def forward(self, x, labels=None):
        B, seq_len = x.size()
        if seq_len > SEQUENCE_LENGTH:
            raise ValueError(f"Input sequence length {seq_len} exceeds maximum {SEQUENCE_LENGTH}")
        x = self.embed(x)
        global_token = self.global_token.expand(B, -1, -1)
        x = torch.cat([global_token, x], dim=1)
        x = x + self.pos_embed[:, :seq_len + 1, :]

        x = self.memory_module(x)

        x1, x2 = x.chunk(2, dim=-1)
        for layer in self.layers:
            x1, x2 = layer(x1, x2)
        x = torch.cat([x1, x2], dim=-1)
        x = self.norm(x)
        x = x[:, 1:, :]
        logits = self.out(x)

        if labels is not None:
            confidence_scores, confidence_loss = self.confidence_predictor(x, labels)
            return logits, confidence_loss
        else:
            return logits

class EfficientOpenWebTextDataset(Dataset):
    def __init__(self, tokenizer, sequence_length):
        self.dataset = load_dataset("wikipedia", "20220301.en", split="train")
        self.tokenizer = tokenizer
        self.sequence_length = sequence_length

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        text = self.dataset[idx]["text"]
        encoding = self.tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=self.sequence_length,
            return_tensors="pt"
        )
        input_ids = encoding["input_ids"].squeeze()
        return {
            "input_ids": input_ids,
            "labels": input_ids.clone()
        }

def train(model, dataloader, optimizer, num_epochs):
    model.train()
    scaler = torch.cuda.amp.GradScaler()
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id, reduction='none')

    for epoch in range(num_epochs):
        print(f"Epoch {epoch + 1}")
        total_loss = 0
        total_confidence_loss = 0
        total_steps = 0
        start_time = time.time()
        pbar = tqdm(enumerate(dataloader), total=len(dataloader))
        for step, batch in pbar:
            input_ids = batch["input_ids"].to(device)
            labels = batch["labels"].to(device)

            optimizer.zero_grad()
            with torch.cuda.amp.autocast():
                outputs, confidence_loss = model(input_ids, labels)
                shift_logits = outputs[..., :-1, :].contiguous()
                shift_labels = labels[..., 1:].contiguous()
                
                # Reshape logits and labels
                shift_logits = shift_logits.view(-1, shift_logits.size(-1))
                shift_labels = shift_labels.view(-1)
                
                # Calculate loss
                loss = criterion(shift_logits, shift_labels)
                
                # Create mask
                mask = (shift_labels != tokenizer.pad_token_id).float()
                
                # Apply mask and calculate mean
                loss = (loss * mask).sum() / mask.sum()
                
                total_loss_step = loss + confidence_loss

            if torch.isnan(total_loss_step).any() or torch.isinf(total_loss_step).any():
                print(f"NaN or Inf loss detected at step {step}. Skipping this batch.")
                continue

            scaler.scale(total_loss_step).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()

            total_loss += loss.item()
            total_confidence_loss += confidence_loss.item()
            total_steps += 1

            if (step + 1) % SAVE_STEPS == 0:
                save_and_infer(model, step)

            iter_per_sec = (step + 1) / (time.time() - start_time)
            avg_loss = total_loss / total_steps
            avg_conf_loss = total_confidence_loss / total_steps
            pbar.set_description(f"Loss: {loss.item():.4f}, Conf Loss: {confidence_loss.item():.4f}, Avg Loss: {avg_loss:.4f}, Iter/s: {iter_per_sec:.2f}")

        for name, param in model.named_parameters():
            if param.grad is not None:
                if torch.isnan(param.grad).any() or torch.isinf(param.grad).any():
                    print(f"NaN or Inf gradient detected in {name}")

    return total_loss / total_steps

def save_and_infer(model, step):
    torch.save(model.state_dict(), MODEL_PATH)
    print(f"Model saved at step {step + 1}")

    model.eval()
    prompt = "The capital of France is"
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)

    max_length = 50
    generated = input_ids

    with torch.no_grad():
        for _ in range(max_length - input_ids.size(1)):
            outputs = model(generated)
            next_token_logits = outputs[:, -1, :] / 0.7
            filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=50, top_p=0.9)
            next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)
            generated = torch.cat([generated, next_token], dim=1)
            if next_token.item() == tokenizer.eos_token_id or generated.size(1) >= SEQUENCE_LENGTH:
                break

    generated_text = tokenizer.decode(generated[0], skip_special_tokens=True)
    print(f"Generated text: {generated_text}")

    model.train()

def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering."""
    logits = logits.clone()
    top_k = min(top_k, logits.size(-1))  # Safety check

    if top_k > 0:
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value

    if top_p > 0.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1)

        # Remove tokens with cumulative probability above top_p
        sorted_indices_to_remove = cumulative_probs > top_p

        # Shift the indices to include the first token surpassing top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        for i in range(logits.size(0)):
            indices_to_remove = sorted_indices[i][sorted_indices_to_remove[i]]
            logits[i, indices_to_remove] = filter_value

    return logits

def load_or_create_model():
    vocab_size = tokenizer.vocab_size
    if os.path.exists(MODEL_PATH):
        print(f"Loading existing model from {MODEL_PATH}")
        model = ReversibleTransformer(NUM_LAYERS, EMBEDDING_DIM, HEADS, vocab_size, MEMORY_SIZE)
        model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
    else:
        print("Creating new model")
        model = ReversibleTransformer(NUM_LAYERS, EMBEDDING_DIM, HEADS, vocab_size, MEMORY_SIZE)
    
    print(f"Model created with {NUM_LAYERS} layers, {EMBEDDING_DIM} embedding dim, and {HEADS} heads")
    return model.to(device)

def main():
    print("Initializing dataset and dataloader...")
    dataset = EfficientOpenWebTextDataset(tokenizer, SEQUENCE_LENGTH)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)

    print("Loading or creating model...")
    model = load_or_create_model()

    num_params = count_parameters(model)
    print(f"Number of trainable parameters: {num_params:,}")

    print("Initializing optimizer...")
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, betas=(0.9, 0.95), eps=1e-8)

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5, verbose=True)

    print("Starting initial inference...")
    save_and_infer(model, step=0)

    print("Starting training...")
    for epoch in range(NUM_EPOCHS):
        avg_loss = train(model, dataloader, optimizer, 1)
        print(f"Epoch {epoch + 1} completed. Average loss: {avg_loss:.4f}")
        scheduler.step(avg_loss)

    print("Performing final save and inference...")
    save_and_infer(model, step=-1)

    print("Training and evaluation complete.")

if __name__ == "__main__":
    main()