CH05- got weird result from train_model_simple

[Nevermetyou65]

Hi
I am reading ch5 pre-training LLM
When I tried to run train_model_simple function here is the result I got

Ep 1 (Step 000000): Train loss 10.858, Val loss 10.893
Ep 1 (Step 000005): Train loss 9.957, Val loss 10.081
Every effort moves you for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for for

Ep 2 (Step 000010): Train loss 8.829, Val loss 9.059
Ep 2 (Step 000015): Train loss 7.773, Val loss 8.063
Every effort moves you;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

Ep 3 (Step 000020): Train loss 7.052, Val loss 7.390
Ep 3 (Step 000025): Train loss 6.690, Val loss 7.101
Every effort moves you..................................................

Ep 4 (Step 000030): Train loss 6.445, Val loss 6.859
Ep 4 (Step 000035): Train loss 6.355, Val loss 6.758
Every effort moves you..................................................

Ep 5 (Step 000040): Train loss 6.787, Val loss 7.138
Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,

Ep 6 (Step 000045): Train loss 6.139, Val loss 6.646
Ep 6 (Step 000050): Train loss 6.124, Val loss 6.740
Every effort moves you the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the

Ep 7 (Step 000055): Train loss 6.287, Val loss 6.747
Ep 7 (Step 000060): Train loss 6.090, Val loss 6.691
Every effort moves you                                                  

Ep 8 (Step 000065): Train loss 6.131, Val loss 6.954
Ep 8 (Step 000070): Train loss 6.088, Val loss 6.800
Every effort moves you,.................................................

Ep 9 (Step 000075): Train loss 6.423, Val loss 6.919
Ep 9 (Step 000080): Train loss 6.423, Val loss 6.894
Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,

Ep 10 (Step 000085): Train loss 6.034, Val loss 6.733
Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,

Which is kind of different from that printed in the book.
Can someone help me?

Here is my code

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)

num_epochs = 10
train_losses, val_losses, tokens_seen = train_model_simple(
    model, train_loader, val_loader, optimizer, device,
    num_epochs=num_epochs, eval_freq=5, eval_iter=5,
    start_context="Every effort moves you", tokenizer=tokenizer
)
def train_model_simple(
    model,
    train_loader,
    val_loader,
    optimizer,
    device,
    num_epochs,
    eval_freq,
    eval_iter,
    start_context,
    tokenizer,
):
    # Initialize lists to track losses and tokens seen
    train_losses, val_losses, track_tokens_seen = [], [], []
    tokens_seen, global_step = 0, -1

    # Main training loop
    for epoch in range(num_epochs):
        model.train()  # Set model to training mode

        for input_batch, target_batch in train_loader:
            optimizer.zero_grad()  # Reset loss gradients from previous batch iteration
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            loss.backward()  # Calculate loss gradients
            optimizer.step()  # Update model weights using loss gradients
            tokens_seen += input_batch.numel()
            global_step += 1

            # Optional evaluation step
            if global_step % eval_freq == 0:
                train_loss, val_loss = evaluate_model(
                    model, train_loader, val_loader, device, eval_iter
                )
                train_losses.append(train_loss)
                val_losses.append(val_loss)
                track_tokens_seen.append(tokens_seen)
                print(
                    f"Ep {epoch+1} (Step {global_step:06d}): "
                    f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}"
                )

        # Print a sample text after each epoch
        generate_and_print_sample(model, tokenizer, device, start_context)
        print()

    return train_losses, val_losses, track_tokens_seen

[rasbt]

Hi there,

thanks for sharing, and hm, that's weird. The code looks correct but yeah, based on the high training loss, something is not right and the model is not learning. Could be that there's an issue with the dataset, or maybe the model configuration itself. Maybe the architecture is not set up correctly.

For the remaining code where, can you use the functions from here:

https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/01_main-chapter-code/previous_chapters.py

It's just so that we can narrow down where the typo occurs.

[Nevermetyou65]

Yes, I am using torch 2.4.1.

[rasbt]

I see. Hm, that's a bit worse than I expected but yeah, it's very likely that it's due to the MPS support in PyTorch since CPU trains normally in your case, correct?

[Nevermetyou65]

Yes, that's correct. Don't worry, sir; I can just use the CPU for now.

[rasbt]

Sounds good! Btw I also have recommendations for GPU cloud resources here if it helps: https://github.com/rasbt/LLMs-from-scratch/tree/main/setup#cloud-resources

As a general rule, even though I also use a Mac for my personal projects and my day job, I wouldn't use it for model training because of the problems with MPS in PyTorch. Inference is a bit better but yeah for training I would not recommend it yet.

[Nevermetyou65]

Yes, Thanks for the advice, sir.