Support mistralai/Mistral-Nemo-Base-2407

🚀 Model / language coverage

Support the Mistral-Nemo-Base-2407 model in Thunder.

• #1279
• https://github.com/NVIDIA/Fuser/issues/3227

Pitch

This is an ask from the NeMo team.

Minimal Repro

Current version of the model:

import torch
from torch.utils.data import DataLoader
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, get_scheduler
from transformers import get_scheduler
from huggingface_hub import snapshot_download
from pathlib import Path
from datasets import load_dataset
import thunder
import thunder.dynamo

mistral_models_path = Path.home().joinpath('mistral_models', 'Nemo-v0.1')
mistral_models_path.mkdir(parents=True, exist_ok=True)

snapshot_download(
  repo_id="mistralai/Mistral-Nemo-Base-2407",
  allow_patterns=["params.json", "consolidated.safetensors", "tekken.json"],
  local_dir=mistral_models_path
)

model_id = "mistralai/Mistral-Nemo-Base-2407"
tokenizer = AutoTokenizer.from_pretrained(
  model_id,
  torch_dtype=torch.bfloat16,
  ignore_mismatched_sizes=True,
)

config = AutoConfig.from_pretrained(model_id)
config.num_hidden_layers = 2
config.torch_dtype = torch.bfloat16
config.ignore_mismatched_sizes=True
config.max_position_embeddings=1024
model = AutoModelForCausalLM.from_config(config)
model = torch.compile(model, backend=thunder.dynamo.ThunderCompiler())
#model = torch.compile(model, backend='eager')

# Add a padding token to the tokenizer
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
    model.resize_token_embeddings(len(tokenizer))

#inputs = tokenizer("Hello my name is", return_tensors="pt")
#outputs = model.generate(**inputs, max_new_tokens=20)
#print(tokenizer.decode(outputs[0], skip_special_tokens=True))

dataset = load_dataset("tiny_shakespeare", split='train',
                       trust_remote_code=True)

# Tokenize the dataset
def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True,
                     max_length=2)

tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# Convert the dataset to PyTorch format and specify columns to return as tensors
tokenized_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask'])

# Create PyTorch DataLoader
dataloader = DataLoader(tokenized_dataset, batch_size=1, shuffle=True)

# Define optimizer and learning rate scheduler
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
num_epochs = 3
lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_epochs * len(dataloader),
)

# Move model to GPU if available
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)

# Fine-tuning loop
model.train()
for epoch in range(num_epochs):
    total_loss = 0
    for batch in dataloader:
        # Move input tensors to device
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)

        # Forward pass
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
        loss = outputs.loss
        print(loss)
        total_loss += loss.item()

        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Update learning rate
        lr_scheduler.step()

    avg_loss = total_loss / len(dataloader)
    print(f"Epoch {epoch + 1}/{num_epochs} completed. Average Loss: {avg_loss:.4f}")

bytes = torch.cuda.memory.max_memory_allocated()
mb: int = 1024*1024
gb: int = 1024*1024*1024
print(f"max allocated: {bytes/1024}kB {bytes/mb}mB {bytes/gb}gB")

cc @tfogal

Start by running the provided minimal reproduction with AutoModelForCausalLM and ThunderCompiler, focusing on the Mistral-Nemo-Base-2407 configuration and training loop. Trace the failure through the Thunder Dynamo integration and compare it with the linked issues. Done means the model's reduced training example runs successfully under Thunder.