GPT 训练管线全流程
GPT 训练管线涉及数据预处理、Tokenizer 训练、模型配置、训练循环和检查点管理等多个环节。
训练循环
python
class GPTTrainer:
def __init__(self, model, optimizer, scheduler, grad_accum_steps=8):
self.model = model
self.optimizer = optimizer
self.scheduler = scheduler
self.grad_accum_steps = grad_accum_steps
def train_step(self, batch):
input_ids, labels = batch["input_ids"], batch["labels"]
with torch.cuda.amp.autocast():
logits = self.model(input_ids)
loss = F.cross_entropy(
logits.view(-1, logits.size(-1)), labels.view(-1),
ignore_index=-100
) / self.grad_accum_steps
loss.backward()
return loss.item() * self.grad_accum_steps
def train_epoch(self, dataloader):
self.model.train()
total_loss = 0
for i, batch in enumerate(dataloader):
loss = self.train_step(batch)
total_loss += loss
if (i + 1) % self.grad_accum_steps == 0:
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.optimizer.step()
self.scheduler.step()
self.optimizer.zero_grad()
return total_loss / len(dataloader)梯度累积
梯度累积模拟更大的 batch size——每 N 步才执行一次优化器更新。在 GPU 内存有限时，这是训练大模型的关键技巧。