the loss is decrasing very fast and you try and find what might be wrong in this llm training code ? 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 = 12
SEQUENCE_LENGTH = 256
LEARNING_RATE = 2e-5
NUM_EPOCHS = 3
SAVE_STEPS = 500
EMBEDDING_DIM = 4096
NUM_LAYERS = 12
HEADS = 16
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 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_qk = nn.Linear(embed_dim, embed_dim * 2)
        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_qk.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_qk.bias)
        nn.init.zeros_(self.to_v.bias)
        nn.init.zeros_(self.to_out.bias)

    def forward(self, x):
        x = self.norm(x)
        B, T, E = x.size()
        qk = self.to_qk(x).view(B, T, self.heads, 2 * self.head_dim)
        q, k = qk.chunk(2, dim=-1)
        v = self.to_v(x).view(B, T, self.heads, self.head_dim)

        q = F.softmax(q, dim=-2)
        k = F.softmax(k, dim=-2)

        # Create causal mask
        mask = torch.tril(torch.ones(T, T)).unsqueeze(0).unsqueeze(0).to(x.device)

        context = torch.einsum('bthd,bshd->bhts', k, v)
        context = context * mask
        out = torch.einsum('bhts,bthd->bshd', context, q)
        out = out.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 = CausalLinearAttention(embed_dim_half, heads)
        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 ReversibleTransformer(nn.Module):
    def __init__(self, depth, embed_dim, heads, vocab_size):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.pos_embed = nn.Parameter(torch.zeros(1, SEQUENCE_LENGTH, 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.init_weights()

    def init_weights(self):
        nn.init.normal_(self.embed.weight, std=0.02)
        nn.init.normal_(self.pos_embed, 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):
        seq_len = x.size(1)
        if seq_len > SEQUENCE_LENGTH:
            raise ValueError(f"Input sequence length {seq_len} exceeds maximum {SEQUENCE_LENGTH}")
        x = self.embed(x) + self.pos_embed[:, :seq_len, :]
        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)
        return self.out(x)

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 + 1,
            return_tensors="pt"
        )
        input_ids = encoding["input_ids"].squeeze()
        return {
            "input_ids": input_ids[:self.sequence_length],
            "labels": input_ids[1:self.sequence_length+1]
        }

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

    for epoch in range(num_epochs):
        print(f"Epoch {epoch + 1}")
        total_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.autocast('cuda'):
                outputs = model(input_ids)
                # Shift the outputs and labels for causal language modeling
                shift_logits = outputs[..., :-1, :].contiguous()
                shift_labels = labels[..., 1:].contiguous()
                loss = criterion(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

            if torch.isnan(loss).any():
                print(f"NaN loss detected at step {step}. Input shape: {input_ids.shape}, Output shape: {outputs.shape}")
                print(f"Input sample: {input_ids[0][:10]}")
                print(f"Output sample: {outputs[0][0][:10]}")
                continue

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

            total_loss += 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
            pbar.set_description(f"Loss: {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():
                    print(f"NaN 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[0, -1, :]
            next_token_logits = next_token_logits / 0.7
            top_k = 50
            top_k_logits, top_k_indices = torch.topk(next_token_logits, top_k)
            top_k_probs = F.softmax(top_k_logits, dim=-1)
            next_token_index = torch.multinomial(top_k_probs, num_samples=1)
            next_token = top_k_indices[next_token_index]
            generated = torch.cat([generated, next_token.unsqueeze(0)], 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 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)
        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)
    
    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__":
    try:
        main()
    except Exception as e:
        print(f"An error occurred: {e}")