Why is checkpointing important?
Checkpointing is a common practice in LLM training due to the long duration and resource-intensive nature of the LLM training process.
Let’s implement a basic checkpointing system:
import torch
from transformers import GPT2LMHeadModel, GPT2Config
import os
class LLMTrainer:
def __init__(
self, model, optimizer, checkpoint_dir='checkpoints'
):
self.model = model
self.optimizer = optimizer
self.checkpoint_dir = checkpoint_dir
os.makedirs(checkpoint_dir, exist_ok=True)
def save_checkpoint(self, epoch, step, loss):
checkpoint = {
...
