Zero-to-GPU / 零基础上云训练
从安装 Python 包到训练 MNIST MLP
这份教程面向第一次使用 DuanFlow 的老师和学生。你会看到一份普通 PyTorch MNIST MLP 训练代码,其中 黄色高亮 的几行就是 DuanFlow 相关代码;其它部分和在本地电脑或学校机房里跑 PyTorch 没区别。
核心思路
普通代码照旧写
模型、数据集、训练循环、loss、optimizer 都是标准 PyTorch。老师原来的教学代码不用重写。
只加高亮几行
高亮的 DuanFlow 代码负责声明镜像、选择 GPU/NPU,并把函数提交到服务器运行。
看代码时只记住一句话:
没有高亮的部分 = 本地 PyTorch 训练;黄色高亮的部分 = 把这段训练放到 DuanFlow 算力服务器上跑。
1. 安装 DuanFlow Python 包
在自己的电脑、学校实验室电脑或教师机上执行。这里安装的是提交任务的 SDK,不是在本机训练大模型。
terminal
python -m pip install https://duanflow.zeabur.app/sdk/duanflow-0.1.0.tar.gz
2. 浏览器授权登录
学校可以给每个老师、学生开账号。登录后,SDK 会带着你的身份提交任务,后台就能统计谁用了多少小时、产生多少费用。不要把密码、API key 或 token 粘贴给 Agent。
terminal
duanflow auth login --api-base https://duanflow.zeabur.app
duanflow whoami
3. 范例 A:使用昇腾 Ascend 910B 训练 MNIST MLP
下面代码中,黄色高亮的行是 DuanFlow 相关代码;其余就是标准 PyTorch 训练代码。昇腾 910B 使用 torch-npu,所以镜像依赖里写了 torch 和 torch-npu。
mnist_ascend_910b.py
import duanflow as df # DuanFlow:引入平台 SDK
image = df.Image.from_registry("registry.cmcc.example.com/synapse/ascend-runtime:latest").uv_pip_install("torch", "torch-npu", "numpy") # DuanFlow:声明服务器镜像和依赖
app = df.App("school-mnist-ascend", image=image) # DuanFlow:创建一个云端应用
@app.function(gpu="Ascend 910B", gpu_count=1, timeout=1800) # DuanFlow:指定使用 1 张昇腾 910B
def train_mnist():
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torch_npu # 昇腾环境需要;如果在本地 CPU 跑,可以删除这一行
device = torch.device("npu:0") # 普通 PyTorch 写法,只是设备从 cuda 换成 npu
# 下载并标准化 MNIST。这里和本地训练完全一样。
tfm = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_ds = datasets.MNIST("/tmp/mnist", train=True, download=True, transform=tfm)
test_ds = datasets.MNIST("/tmp/mnist", train=False, download=True, transform=tfm)
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=1024)
# 一个最简单的 MLP:784 -> 256 -> 10
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
# 训练 3 个 epoch。这里仍然是普通 PyTorch。
for epoch in range(3):
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
loss = F.cross_entropy(model(x), y)
loss.backward()
opt.step()
print(f"epoch={epoch + 1} loss={loss.item():.4f}")
# 测试准确率。
model.eval()
correct = total = 0
with torch.no_grad():
for x, y in test_loader:
pred = model(x.to(device)).argmax(1).cpu()
correct += (pred == y).sum().item()
total += y.numel()
return {"accuracy": correct / total}
print(train_mnist.remote()) # DuanFlow:把函数提交到服务器运行;本地只负责发起任务
4. 范例 B:使用 NVIDIA RTX 4090 训练 MNIST MLP
如果学校或平台接入的是 NVIDIA 4090,代码几乎一样,只是镜像和 GPU 名称不同,设备使用 cuda。
mnist_4090.py
import duanflow as df # DuanFlow:引入平台 SDK
image = df.Image.debian_slim().uv_pip_install("torch", "torchvision", "numpy") # DuanFlow:声明 Python 依赖
app = df.App("school-mnist-4090", image=image) # DuanFlow:创建一个云端应用
@app.function(gpu="RTX 4090", gpu_count=1, timeout=1800) # DuanFlow:指定使用 1 张 NVIDIA 4090
def train_mnist():
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 本地 PyTorch 常见写法
tfm = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_ds = datasets.MNIST("/tmp/mnist", train=True, download=True, transform=tfm)
test_ds = datasets.MNIST("/tmp/mnist", train=False, download=True, transform=tfm)
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=1024)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
for epoch in range(3):
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
loss = F.cross_entropy(model(x), y)
loss.backward()
opt.step()
print(f"epoch={epoch + 1} loss={loss.item():.4f}")
model.eval()
correct = total = 0
with torch.no_grad():
for x, y in test_loader:
pred = model(x.to(device)).argmax(1).cpu()
correct += (pred == y).sum().item()
total += y.numel()
return {"accuracy": correct / total}
print(train_mnist.remote()) # DuanFlow:提交到服务器;结果和日志回到本地
5. 运行与查看日志
运行方式和普通 Python 文件一样。区别是:训练不会占用本地电脑 GPU,而是提交到学校后台配置的 DuanFlow 算力服务器。
terminal
python mnist_ascend_910b.py
# 或者
python mnist_4090.py
学生看到什么终端会打印训练日志和最终 accuracy。学校后台会记录这次任务。
老师看到什么老师能在学校后台看到谁跑了任务、用了多少卡时、是否成功、日志是什么。