从零手搓大模型前置知识(附录二)PyTorch GPU 训练基础

发布时间:2026/7/4 20:08:27
从零手搓大模型前置知识(附录二)PyTorch GPU 训练基础 从零手搓大模型前置知识附录二PyTorch GPU 训练基础接着附录一重点讲如何把 tensor、模型、数据移动到 GPU 上训练如果你没有 NVIDIA GPU也可以阅读这篇因为它解释的是后面训练大模型必须理解的设备管理逻辑。1. GPU 训练为什么重要深度学习训练里大量计算都是矩阵乘法。GPU 非常擅长并行矩阵计算所以训练神经网络时通常比 CPU 快很多。但是 PyTorch 有一个重要规则参与同一次计算的 tensor 必须在同一个设备上。也就是说不能让一个 tensor 在 CPU另一个 tensor 在 GPU然后直接相加或矩阵乘法。2. 检查 PyTorch 和 GPUimporttorchprint(torch.__version__)检查 CUDA 是否可用print(torch.cuda.is_available())如果输出True说明 PyTorch 可以使用 NVIDIA GPU。如果是False说明当前环境没有可用 CUDA。原因可能是没有 NVIDIA 显卡。没安装 CUDA 版 PyTorch。驱动或环境没配好。3. CPU tensor 运算先创建两个 tensortensor_1torch.tensor([1.,2.,3.])tensor_2torch.tensor([4.,5.,6.])print(tensor_1.device)print(tensor_1tensor_2)输出默认情况下它们在 CPU 上。可以查看print(tensor_1.device)一般是cpu4. 把 tensor 移动到 GPU如果 CUDA 可用可以写tensor_1tensor_1.to(cuda)tensor_2tensor_2.to(cuda)print(tensor_1tensor_2)输出.to(cuda)的意思是把 tensor 复制/移动到 GPU 设备上也可以移回 CPUtensor_1tensor_1.to(cpu)5. 最常见的设备错误如果一个 tensor 在 CPU另一个在 GPUtensor_1tensor_1.to(cpu)print(tensor_1tensor_2)这通常会报错。原因是tensor_1 在 CPU tensor_2 在 CUDAPyTorch 不允许它们直接相加。所以训练时要记住模型、输入、标签必须在同一个 device 上。6. 更稳的设备写法真实代码里一般不直接写死cuda而是devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)这样有 GPU 就用 GPU 没有 GPU 就自动退回 CPU后面所有东西都往这个device上移动model.to(device)featuresfeatures.to(device)labelslabels.to(device)7. 准备训练数据复用了附录一的玩具数据X_traintorch.tensor([[-1.2,3.1],[-0.9,2.9],[-0.5,2.6],[2.3,-1.1],[2.7,-1.5]])y_traintorch.tensor([0,0,0,1,1])X_testtorch.tensor([[-0.8,2.8],[2.6,-1.6],])y_testtorch.tensor([0,1])这是一个简单二分类问题。8. Dataset 仍然一样Dataset 和 CPU 训练时一样fromtorch.utils.dataimportDatasetclassToyDataset(Dataset):def__init__(self,X,y):self.featuresX self.labelsydef__getitem__(self,index):one_xself.features[index]one_yself.labels[index]returnone_x,one_ydef__len__(self):returnself.labels.shape[0]train_dsToyDataset(X_train,y_train)test_dsToyDataset(X_test,y_test)注意通常 Dataset 里的原始数据可以先放 CPU。训练循环里每取出一个 batch再把 batch 移动到 GPU。9. DataLoaderfromtorch.utils.dataimportDataLoader torch.manual_seed(123)train_loaderDataLoader(datasettrain_ds,batch_size2,shuffleTrue,num_workers1,drop_lastTrue)test_loaderDataLoader(datasettest_ds,batch_size2,shuffleFalse,num_workers1)参数drop_lastTrue的意思是如果最后一个 batch 不够 batch_size就丢掉。在某些训练场景里这可以让每个 batch 形状保持一致。Windows 环境如果遇到 DataLoader 多进程问题可以把num_workers1改成num_workers010. 定义模型模型和之前的附录一是一样的classNeuralNetwork(torch.nn.Module):def__init__(self,num_inputs,num_outputs):super().__init__()self.layerstorch.nn.Sequential(torch.nn.Linear(num_inputs,30),torch.nn.ReLU(),torch.nn.Linear(30,20),torch.nn.ReLU(),torch.nn.Linear(20,num_outputs),)defforward(self,x):logitsself.layers(x)returnlogitsCPU 训练和 GPU 训练的模型定义通常不需要不同。不同的是模型创建后要移动到 device。11. 单 GPU 训练循环这是这部分最重要的代码importtorch.nn.functionalasF torch.manual_seed(123)modelNeuralNetwork(num_inputs2,num_outputs2)devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)model.to(device)optimizertorch.optim.SGD(model.parameters(),lr0.5)num_epochs3forepochinrange(num_epochs):model.train()forbatch_idx,(features,labels)inenumerate(train_loader):features,labelsfeatures.to(device),labels.to(device)logitsmodel(features)lossF.cross_entropy(logits,labels)optimizer.zero_grad()loss.backward()optimizer.step()print(fEpoch:{epoch1:03d}/{num_epochs:03d}f | Batch{batch_idx1:03d}/{len(train_loader):03d}f | Train/Val Loss:{loss:.2f})model.eval()DataLoader 迭代只是批量调度器真正读取单条数据的核心入口永远是 Dataset 的 getitem。相比附录一只多了几行关键代码。创建 devicedevicetorch.device(cudaiftorch.cuda.is_available()elsecpu)移动模型model.to(device)移动每个 batchfeatures,labelsfeatures.to(device),labels.to(device)这是 GPU 训练最核心的三步。12. 为什么每个 batch 都要.to(device)DataLoader 每次取出来的features和labels默认通常在 CPU。而模型已经被移动到 GPUmodel.to(device)所以输入也必须移动到同一个设备featuresfeatures.to(device)labelslabels.to(device)否则你会看到类似错误Expected all tensors to be on the same device这个错误非常常见。13. GPU 版本的准确率函数现在对compute_accuracy也做了设备适配defcompute_accuracy(model,dataloader,device):modelmodel.eval()correct0.0total_examples0foridx,(features,labels)inenumerate(dataloader):features,labelsfeatures.to(device),labels.to(device)withtorch.no_grad():logitsmodel(features)predictionstorch.argmax(logits,dim1)comparelabelspredictions correcttorch.sum(compare)total_exampleslen(compare)return(correct/total_examples).item()和 CPU 版相比关键新增features,labelsfeatures.to(device),labels.to(device)评估时也要保证模型、输入、标签在同一个 device 上。14. 训练后评估训练集准确率compute_accuracy(model,train_loader,devicedevice)测试集准确率compute_accuracy(model,test_loader,devicedevice)这里的devicedevice是显式传入当前设备。这样函数里就知道应该把 batch 移到哪里。15. 多 GPU 训练DDP 是 Distributed Data Parallel分布式数据并行。初学阶段不用急着看。建议顺序是先会 CPU 训练 再会单 GPU 训练 最后再看多 GPU / DDP训练大模型时多 GPU 很重要但理解手搓 LLM 主线时单 GPU 思维已经够用了。16.本章和 LLM 的关系后面训练 GPT 时也会用同样逻辑devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)model.to(device)forinput_batch,target_batchintrain_loader:input_batchinput_batch.to(device)target_batchtarget_batch.to(device)logitsmodel(input_batch)lossloss_fn(logits,target_batch)optimizer.zero_grad()loss.backward()optimizer.step()也就是说Part 2 学的是训练大模型前必须养成的设备习惯。17. 你最该记住的模板单 GPU/CPU 自适应训练模板devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)modelNeuralNetwork(num_inputs2,num_outputs2)model.to(device)forfeatures,labelsintrain_loader:featuresfeatures.to(device)labelslabels.to(device)logitsmodel(features)lossF.cross_entropy(logits,labels)optimizer.zero_grad()loss.backward()optimizer.step()这套模板后面会反复出现只是模型从小 MLP 换成 GPT数据从二维点换成 token 序列。18. 常见错误排查错误 1CUDA 不可用torch.cuda.is_available()输出False就不要强行.to(cuda)。用devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)错误 2模型在 GPU数据在 CPU解决featuresfeatures.to(device)labelslabels.to(device)错误 3评估时忘记移动数据评估函数里也要写features,labelsfeatures.to(device),labels.to(device)错误 4Windows DataLoader 报多进程问题把num_workers1改成num_workers019. 建议学习顺序先运行torch.cuda.is_available()确认有没有 GPU。如果没有 GPU也继续看.to(device)逻辑。跑 tensor.to(cuda)示例时如果没 GPU跳过硬编码cuda的单元。重点跑单 GPU 训练循环把model.to(device)和features.to(device)记牢。多 GPU DDP 先跳过等主线学完再回来看。一句话总结附录一教你怎么训练模型附录二教你怎么把训练搬到 GPU 上。