1. 这不是数学课是训练大模型的“方向盘”和“刹车系统”你刚打开一个大模型训练脚本optimizer.step()执行完loss 下降了 0.002——但你真的知道这行代码背后发生了什么吗它不是魔法也不是黑箱里自动吐出的结果。它是一整套精密协作的机械梯度下降是方向盘决定模型往哪走反向传播是刹车油门后视镜的组合体告诉你当前方向对不对、该踩多深、上一步哪里偏了mini_batch 是你每次只看一小段路既省油又防晕车而计算图就是你脑中那张实时更新的导航地图——没有它你连自己在哪、要去哪、怎么调头都不知道。这不是抽象理论而是每天在 GPU 显存里真实发生的物理过程。我带过三轮大模型训练项目从 7B 参数量的 LLaMA 微调到自研 MoE 架构的千卡集群训练最常被问的问题从来不是“怎么搭环境”而是“为什么 loss 突然炸了”“为什么梯度全为零”“为什么 batch_size 改成 8 就 OOM改成 4 却训不动”——这些问题的答案90% 都藏在这四个词里梯度下降、反向传播、mini_batch、计算图。它们不是独立模块而是一个闭环系统。你调参时改 learning_rate是在动方向盘的灵敏度你加 gradient clipping是在给刹车加助力泵你换torch.compile()是在重绘那张计算图的拓扑结构你用torch.utils.data.DataLoader设置drop_lastTrue是在确保每一段“小路”长度一致避免导航失准。这篇文章不讲推导证明不列拉格朗日乘子不画 sigmoid 函数求导链式法则。我们只做一件事把这四个概念还原成你在终端里敲命令、在 PyTorch 里写loss.backward()、在 TensorBoard 里看曲线时真正能感知、能干预、能 debug 的具体动作。你会看到梯度下降不是“沿着坡往下滚”而是“每步都重新测绘坡度再决定落脚点”的动态勘测过程反向传播不是“从输出倒着算导数”而是计算图上一场有严格时序、带内存地址追踪的“信号回溯风暴”mini_batch 不是“把数据切小块”而是训练稳定性的核心调节阀它的大小直接决定你能否用上 8 张 A100 而不触发 NCCL timeout计算图不是静态 DAG 图而是 PyTorch Autograd 引擎在每次 forward 时现场生成、带唯一 ID、可被torch.autograd.grad()显式干预的活体结构。如果你正卡在 loss 不降、显存爆满、梯度消失/爆炸、multi-GPU 同步失败这些高频问题上那么你缺的不是新模型而是对这四个基础机制的“肌肉记忆”。接下来我们就从一次真实的forward → backward → step完整周期出发一帧一帧拆解这个闭环。2. 梯度下降不是下山是每一步都重绘地形图的动态勘测很多人把梯度下降理解成“小球滚下山坡”这个类比在入门阶段有用但一旦进入大模型训练它会成为最大的认知陷阱。真实情况是你不是在已知地形上滚动而是在每一步都用激光雷达扫描局部坡度然后仅凭这一帧扫描结果决定下一步往哪迈、迈多远。地形本身即损失函数曲面在你移动过程中持续变形——因为参数变了模型结构变了甚至数据采样策略也变了。所以梯度下降的本质是一场高维空间里的实时地形测绘 局部决策。2.1 梯度不是标量斜率而是 n 维空间的“指向向量”先破除一个常见误解梯度 ∇L(θ) 不是一个数字而是一个与参数 θ 维度完全相同的向量。假设你训练一个含 10 亿参数的模型那么 ∇L(θ) 就是一个 10^9 维向量每个分量代表对应参数在当前点的偏导数∂L/∂θ_i。它不告诉你“坡有多陡”而是精确指出“在 θ_i 方向上损失函数变化最快的方向和速率”。提示PyTorch 中model.parameters()返回的是一个个Parameter对象每个对象内部.grad属性存储的就是该参数对应的梯度分量。当你执行optimizer.step()时优化器遍历所有.grad按公式θ_i ← θ_i − η × ∂L/∂θ_i更新。这里的关键是η学习率不是全局常量而是每个参数维度上的缩放系数。AdamW 的 weight decay、Layer-wise LR scaling、甚至 LoRA adapter 的独立 lr都是在不同维度上施加不同 η。我们来实测一个 3 层 MLP 的梯度分布。定义模型import torch import torch.nn as nn model nn.Sequential( nn.Linear(100, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 10) )输入一个 batchsize32计算 lossx torch.randn(32, 100) y torch.randint(0, 10, (32,)) loss_fn nn.CrossEntropyLoss() loss loss_fn(model(x), y) loss.backward()此时查看第一层线性层权重的梯度统计w1_grad model[0].weight.grad print(f梯度均值: {w1_grad.mean().item():.6f}) print(f梯度标准差: {w1_grad.std().item():.6f}) print(f梯度最大值: {w1_grad.max().item():.6f}) print(f梯度最小值: {w1_grad.min().item():.6f}) print(f梯度非零比例: {(w1_grad ! 0).float().mean().item():.3%})实测结果典型值梯度均值: -0.000123 梯度标准差: 0.018742 梯度最大值: 0.152341 梯度最小值: -0.148922 梯度非零比例: 100.000%注意标准差 0.0187 远大于均值 0.000123说明梯度整体呈中心对称分布无系统性偏移而最大/最小值接近 ±0.15意味着单个参数更新步长可达 0.15×η。如果你设 learning_rate1e-3那么单步最大更新量是 0.00015 —— 这看起来很小但当参数量达 10^9 时所有梯度分量的累积扰动足以让 loss 曲面剧烈震荡。2.2 学习率不是“调快慢”而是“控步长精度”的校准旋钮学习率 η 的物理意义是控制每次参数更新的绝对步长。它不决定“收敛速度”而决定“能否稳定落在盆地内”。过大则像用大锤敲玻璃——每次更新都越过最优解loss 剧烈震荡甚至发散过小则像用绣花针雕石头——收敛极慢且易陷入尖锐局部极小值大模型中表现为 loss plateau。但关键在于η 必须与梯度幅值匹配。上面实测中梯度 std≈0.0187若设 η1e-2则平均步长 ≈ 0.000187若设 η1e-1则平均步长 ≈ 0.00187是前者的 10 倍——这已超出多数初始化方案的安全范围。我们验证一下# 测试不同 η 下的 loss 变化 for lr in [1e-4, 1e-3, 1e-2, 1e-1]: model.zero_grad() loss loss_fn(model(x), y) loss.backward() for p in model.parameters(): if p.grad is not None: p.data - lr * p.grad print(flr{lr:.0e} - loss{loss_fn(model(x), y).item():.4f})典型输出lr1e-04 - loss2.3142 lr1e-03 - loss2.2876 # 稳定下降 lr1e-02 - loss2.4513 # 开始反弹 lr1e-01 - lossnan # 梯度爆炸loss 变 NaN这就是为什么大模型训练必须用 warmup初始阶段梯度幅值不稳定尤其 embedding 层直接用目标 lr 会一步跨出盆地。warmup 本质是让 η 从 0 缓慢爬升到目标值给计算图和梯度分布一个“热身适应期”。2.3 大模型特有挑战梯度幅值的跨层爆炸与消失在 LLaMA-7B 这样的模型中不同层的梯度幅值差异可达 3 个数量级。我们用 HuggingFace 的transformers加载模型hook 每层梯度def hook_fn(module, grad_input, grad_output): if hasattr(module, weight) and module.weight.grad is not None: print(f{module.__class__.__name__}: grad_std{module.weight.grad.std().item():.3e}) for name, module in model.named_modules(): if isinstance(module, nn.Linear): module.register_backward_hook(hook_fn)实测某次 forward-backward 后各层梯度 std单位e-03embed_tokens.weight: 1.24e-02layers.0.self_attn.q_proj.weight: 8.76e-04layers.10.mlp.gate_proj.weight: 3.12e-05norm.weight: 2.01e-06lm_head.weight: 4.55e-03可见embedding 和 lm_head 层梯度最强中间 transformer 层梯度逐层衰减。这是典型的“梯度消失”现象根源在于 ReLU 或 SiLU 激活函数的导数在负区为 0以及多层矩阵乘法的链式衰减。解决方案不是“加大 lr”而是LayerNorm 位置调整Pre-LN 比 Post-LN 更缓解消失梯度检查点Gradient Checkpointing牺牲 20% 计算时间换取 50% 显存从而允许更大 batchAdaptive Gradient Clipping不设固定阈值而是按层 std 动态 clip如clip_norm 0.1 * layer_grad_std。我在训练 13B 模型时将q_proj层梯度 clip 阈值设为 1.0lm_head设为 0.1loss 曲线从剧烈抖动变为平滑下降——这印证了梯度下降的稳定性取决于你对梯度幅值分布的理解深度而非对公式的背诵熟练度。3. 反向传播一场在计算图上按指令执行的“信号回溯风暴”反向传播Backpropagation常被误认为是“链式法则的自动应用”。错。它是 PyTorch Autograd 引擎在计算图Computational Graph上发起的一场有严格时序、带内存地址追踪、可被中断和重定向的信号回溯风暴。它的核心不是数学而是内存管理 指令调度 依赖解析。3.1 计算图不是静态 DAG而是带生命周期的“活体结构”计算图在 PyTorch 中并非预先构建的静态图如 TensorFlow 1.x而是在每次forward执行时由 Autograd 引擎动态捕获操作并生成的临时结构。每个Tensor都有一个.grad_fn属性指向其创建它的Function对象每个Function又持有输入Tensor的引用和backward方法。整个图是一个以 loss 为根节点、以 leaf tensors如模型参数为叶节点的有向无环图DAG。我们用一个极简例子可视化x torch.tensor(2.0, requires_gradTrue) y x ** 2 z y 3 z.backward() print(fx.grad {x.grad}) # 4.0其计算图为x → yx² → zy3 → lossz。Autograd 在z.backward()时从z开始调用z.grad_fn即AddBackward0它知道z由y和3相加而来于是将梯度1.0因 dz/dz1传给y接着y.grad_fnPowBackward0收到dy/dz1.0计算dx/dy 2x 4最终dx/dz dx/dy * dy/dz 4。关键点在于计算图的生命期与Tensor绑定。一旦y被del y或超出作用域其对应的Function和图节点即被 GC 回收。这也是为什么torch.no_grad()能大幅提速——它直接禁用.grad_fn的创建不生成任何图节点。3.2 反向传播的三阶段准备、回溯、聚合反向传播实际分为三个不可分割的阶段准备阶段Graph Constructionforward执行时Autograd 记录所有可微操作构建grad_fn链。此阶段无计算开销只有内存分配。回溯阶段Backward Passloss.backward()启动从 loss 节点开始按拓扑逆序调用每个Function.backward()计算局部梯度并累加到对应.grad。这是纯 CPU 指令调度不涉及 GPU 计算。聚合阶段Gradient Accumulation当多个 loss如 multi-task共享部分参数时.grad会被多次累加。PyTorch 默认行为是而非覆盖。这是 mini_batch 累积梯度的基础。我们验证第三阶段x torch.tensor(1.0, requires_gradTrue) y1 x ** 2 y2 x ** 3 loss1 y1 loss2 y2 loss1.backward(retain_graphTrue) # retain_graphTrue 允许二次 backward loss2.backward() print(fx.grad {x.grad}) # 2*x 3*x² 2 3 5x.grad是2x来自 loss1和3x²来自 loss2的和。这说明反向传播不是“算一次导数”而是“对所有路径贡献的梯度进行线性叠加”。在大模型中这体现为attention mask 的梯度、label smoothing 的梯度、KL 散度正则项的梯度全部在同一.grad上累加。若某项 loss 权重设得过大会淹没其他信号——这正是 multi-task 训练中 loss balancing 的核心难点。3.3 大模型实战陷阱in-place 操作与图断裂最常导致RuntimeError: Trying to backward through the graph a second time的原因是 in-place 操作破坏了计算图的完整性。例如x torch.tensor(1.0, requires_gradTrue) y x ** 2 y 1 # in-place add修改 y 的内存地址 z y * 2 z.backward() # OK # y.backward() # Error! y 已被 in-place 修改grad_fn 断裂在大模型中这种错误更隐蔽。比如使用F.silu(x)是安全的但x.sigmoid_()in-place sigmoid会破坏图。另一个经典陷阱是torch.cat([a, b], dim0)后直接a[:] ...导致a的梯度无法回传。解决方案只有两个永远用 out-of-place 操作y y 1而非y 1显式启用retain_graphTrue当需多次 backward如 GAN 的 generator/discriminator 交替更新时必须设此 flag否则图在第一次 backward 后即被释放。我在调试一个 MoE 模型时发现 expert routing 的topk操作返回的索引 tensor 若被 in-place 修改会导致后续loss.backward()报错 “leaf variable has been moved into the graph interior”。最终定位到一行indices.clamp_(0, num_experts-1)—— 改为indices torch.clamp(indices, 0, num_experts-1)后问题消失。这再次证明反向传播的健壮性90% 取决于你对 in-place 操作边界的敬畏心。4. mini_batch不是数据切片而是训练稳定性的核心调节阀mini_batch 常被简化为“把大数据集切成小块”。这是严重低估。在大模型训练中mini_batch size 是一个同时影响显存占用、GPU 利用率、梯度噪声水平、分布式同步效率、甚至模型泛化能力的超级参数。它不是“越小越好”或“越大越好”而是一个需要在多个约束间精密平衡的杠杆。4.1 显存消耗的三大组件参数 梯度 激活值一个 batch 的显存占用 模型参数显存 梯度显存 激活值activations显存。其中参数显存固定等于sum(p.numel() for p in model.parameters()) * 2字节FP16梯度显存与参数显存相同因.grad与参数同 dtype、同 shape激活值显存与 batch_size 成正比且随模型深度指数增长。对于 LLaMA-7B单层 attention 的 key/value cache 在 seq_len2048 时约 128MB12 层即超 1.5GB。我们实测不同 batch_size 下的显存占用A100 80GLLaMA-7BFP16batch_size总显存 (GB)参数梯度 (GB)激活值 (GB)128.414.214.2231.614.217.4437.814.223.6849.214.235.0可见batch_size 从 1→8显存增长 73%其中激活值增长 146%。这是因为激活值不仅包括中间 tensor还包括autograd为反向传播保存的 forward 中间结果如 attention softmax 输出。这就是为什么gradient_checkpointing能节省 50% 显存它放弃保存某些中间激活而在 backward 时用recompute重新计算用时间换空间。4.2 梯度噪声batch_size 决定“抽样方差”进而影响收敛路径mini_batch 的本质是用 batch 内样本的梯度均值近似全量数据的梯度期望。根据中心极限定理梯度估计的标准差 σ_grad ∝ 1/√N其中 N 是 batch_size。这意味着N 小 → σ_grad 大 → 梯度方向噪声强 → loss 曲线抖动大但可能跳出局部极小N 大 → σ_grad 小 → 梯度方向精准 → loss 曲线平滑但易陷入尖锐极小泛化性差。我们对比 batch_size1 和 batch_size32 的 loss 曲线同一模型、同一 lrbatch_size1loss 在 2.1~2.9 间剧烈震荡但 1000 step 后降至 1.8batch_size32loss 从 2.5 平滑下降至 2.0但 1000 step 后停滞在 1.95不再下降。这印证了小 batch 提供探索性大 batch 提供收敛性。工业界标准做法是warmup 阶段用小 batch如 1~4快速探路主训练阶段用大 batch如 128~2048稳定收敛finetune 阶段再用中等 batch32~128兼顾泛化。4.3 分布式训练中的 batch_sizeglobal_batch_size 与 micro_batch_size 的分离在千卡集群训练中“batch_size” 一词必须明确是 global 还是 micromicro_batch_size单卡处理的样本数决定单卡显存global_batch_size所有卡累计的样本数决定梯度更新步长。例如128 卡集群micro_batch_size8则 global_batch_size1024。此时每卡计算自己的 loss 和梯度然后通过 all-reduce 同步梯度最后每卡用同步后的梯度更新本地参数。这要求micro_batch_size必须能被单卡显存容纳global_batch_size必须足够大以保证梯度估计的统计可靠性通常 ≥ 2048micro_batch_size过小会导致 all-reduce 通信开销占比过高NCCL 启动延迟显著。我们在 512 卡 A100 上训练 70B 模型时发现当micro_batch_size1时每 step 的 all-reduce 时间占 45%提升至micro_batch_size4后all-reduce 时间占比降至 22%吞吐量提升 2.3 倍。这说明mini_batch size 是连接算法梯度更新与系统GPU/网络的关键接口忽视它就等于用跑车引擎配自行车链条。5. 计算图PyTorch Autograd 的“活体导航地图”可读、可干预、可重绘计算图常被当作黑箱背后的“幕后功臣”。实际上它是 PyTorch 中最透明、最可干预的组件之一。它不是仅供 Autograd 内部使用的隐式结构而是你可以随时 inspect、modify、even replace 的活体对象。理解计算图就是掌握大模型训练的“上帝视角”。5.1 可视化计算图用 torchviz 看清每一帧的拓扑结构安装torchviz后可将任意计算图导出为 DOT 格式并渲染pip install torchvizfrom torchviz import make_dot x torch.tensor(1.0, requires_gradTrue) y x ** 2 z y 3 dot make_dot(z, params{x: x}) dot.render(computational_graph, formatpng, cleanupTrue)生成的图清晰显示x是输入节点PowBackward0是yx²的反向函数AddBackward0是zy3的反向函数z是输出节点。每个节点标注了 tensor shape 和 dtype。这对于 debug 复杂模型如带 condition 的 control flow极为关键。我们曾遇到一个 bug模型在 eval 模式下 loss 为 0train 模式下 loss 爆炸。用make_dot对比发现train 模式下图中多了一个DropoutBackward节点其梯度在 backward 时未被正确归一化。定位到nn.Dropout(p0.1)未设inplaceFalse改为nn.Dropout(p0.1, inplaceFalse)后问题解决。这证明计算图可视化不是炫技而是定位梯度流异常的第一道防线。5.2 干预计算图用 torch.autograd.Function 定义自定义反向逻辑当标准 op 无法满足需求时如量化训练、稀疏更新你需要继承torch.autograd.Function手动定义forward和backwardclass QuantizeLinear(torch.autograd.Function): staticmethod def forward(ctx, input, scale, zero_point): ctx.save_for_backward(input, scale, zero_point) # 量化 forward q_input torch.round(input / scale) zero_point return q_input * scale staticmethod def backward(ctx, grad_output): input, scale, zero_point ctx.saved_tensors # 量化反向直通估计器STE grad_input grad_output.clone() return grad_input, None, None # scale 和 zero_point 不参与梯度更新 # 使用 q_input QuantizeLinear.apply(x, scale, zero_point)这里backward中grad_input grad_output.clone()就是 STE 的核心忽略量化带来的不可微性将梯度“直通”给输入。这在大模型低比特训练中是标配技术。关键点在于ctx.save_for_backward()保存的 tensor在 backward 时可被安全访问且不会增加额外显存——因为它们本就是 forward 中已存在的对象。5.3 重绘计算图torch.compile() 的图优化原理torch.compile()不是简单加速而是对计算图进行多级重绘Level 1Operator Fusion将matmul bias_add relu合并为一个 kernel减少 kernel launch 开销Level 2Memory Layout Optimization重排 tensor 内存布局使连续访存对齐 GPU warpLevel 3Graph Rewriting识别冗余计算如重复softmax插入缓存节点。我们测试 LLaMA-7B 的forward时间原生 PyTorch124ms/steptorch.compile(modedefault)89ms/step-28%torch.compile(modemax-autotune)73ms/step-41%性能提升主要来自图重绘后的 kernel fusion。但要注意compile会改变计算图结构可能导致某些 hook 失效。例如你在Linear层注册的 backward hook在 compile 后可能被融合进更大 kernelhook 不再触发。因此生产环境启用 compile 前必须用torchviz对比编译前后图结构确认关键监控点未被消除。我在部署一个推理服务时启用了max-autotune结果发现梯度裁剪失效——因为clip_grad_norm_作用的.grad被 fuse 进了 optimizer kernel。最终解决方案是在compile前用torch.no_grad()包裹裁剪逻辑确保其独立于主图。这再次强调计算图不是被动容器而是你主动设计、持续维护的训练基础设施。6. 四者闭环一次完整训练 step 的微观世界拆解现在我们将梯度下降、反向传播、mini_batch、计算图放入一次真实的optimizer.step()周期中逐帧拆解这个闭环如何协同工作。以 LLaMA-7B 在 8 卡 A100 上训练为例global_batch_size2048micro_batch_size256learning_rate2e-5。6.1 Step 0DataLoader 加载 mini_batchDataLoader从磁盘读取 256 个 tokenized sequence每个 seq_len2048pad 至统一长度组成input_idstensorshape[256, 2048]。此时数据尚未加载到 GPU仅在 CPU 内存。6.2 Step 1Forward Pass 与计算图生成input_ids input_ids.to(cuda:0) # 拷贝到 GPU outputs model(input_ids) # forward 执行 loss loss_fn(outputs.logits, labels) # 计算 loss在此过程中每个nn.Linear、nn.Embedding、nn.LayerNorm的forward方法被调用Autograd 引擎实时捕获操作为每个输出 tensor 创建grad_fn构建计算图所有中间激活如 attention scores、FFN 输出被保存在 GPU 显存等待 backward图的 root 是losstensorleaf 是model.parameters()。6.3 Step 2Backward Pass 与梯度回溯loss.backward() # 启动反向传播Autograd 引擎从loss节点开始调用CrossEntropyLossBackward依拓扑逆序依次调用LMHeadBackward、TransformerBlockBackward、EmbeddingBackward每个backward方法计算局部梯度并累加到对应参数的.grad此过程纯 CPU 调度GPU 执行的是grad_fn中封装的 CUDA kernel如matmul_backward当所有grad_fn调用完毕所有参数的.grad已被填充。6.4 Step 3Gradient Sync 与 Global Update# DDP 自动执行 # 1. 将所有卡的 .grad 拷贝到 CPU 或专用通信 buffer # 2. 调用 NCCL all-reduce计算 global_grad mean(local_grads) # 3. 将 global_grad 写回每卡参数的 .grad optimizer.step() # 用 global_grad 更新参数 optimizer.zero_grad() # 清空 .grad为下一 batch 准备此时梯度下降完成一次迭代参数 θ 更新为 θ − η × global_grad。而 mini_batch 的使命结束计算图被 GC 回收为下一 batch 的新图腾出空间。6.5 关键洞察四者如何相互制衡计算图的粒度决定了反向传播的路径长度路径越长梯度消失风险越高反向传播的路径长度影响梯度幅值的跨层分布进而要求梯度下降的 lr 分层设置mini_batch size控制激活值显存显存上限又限制了计算图能容纳的最大 seq_len 和 batch_size梯度下降的收敛行为如 loss plateau反过来提示你是否需调整计算图结构如加 checkpoint、或mini_batch 策略如 switch to larger batch。它们不是孤立模块而是一个动态平衡系统。我在调试一个 70B 模型时loss plateau 持续 5000 steps。按常规思路调 lr 无效。最终用torchviz发现最后一层lm_head的grad_fn在 90% 的 steps 中未被调用——定位到labelstensor 的 device 不匹配CPU vs GPU导致 loss 计算跳过反向路径。修复 device 后loss 立即下降。这个案例说明大模型训练的瓶颈往往不在算法前沿而在对这四个基础机制的掌控精度上。7. 实战 checklist上线前必须验证的 7 个硬核指标基于以上分析我整理了一份上线前必须验证的 checklist。它不是理论清单而是我在三次千卡训练中每次部署前必跑的实测脚本。每一条都对应一个真实故障场景。7.1 梯度幅值分布确保无系统性偏移运行 10 个 step收集所有p.grad.std()计算全局 std 的均值与方差grad_stds [] for i in range(10): loss.backward() stds [p.grad.std().item() for p in model.parameters() if p.grad is not None] grad_stds.append(np.mean(stds)) optimizer.zero_grad() print(fgrad_std mean: {np.mean(grad_stds):.3e}, std: {np.std(grad_stds):.3e}) # 合格线mean 1e-4 且 std 0.3 * mean若mean 1e-4说明梯度太小可能初始化不当或激活函数饱和若std/mean 0.3说明梯度分布不稳需检查数据 pipeline 或 loss function。7.2 计算图完整性验证无 in-place 断裂在forward后对每个param检查param.grad_fn是否为Noneleaf tensor 应为None对每个中间 tensor 检查tensor.grad_fn是否非Nonefor name, param in model.named_parameters(): assert param.grad_fn is None, f{name} should be leaf for name, module in model.named_modules(): if hasattr(module, weight) and module.weight.requires_grad: assert module.weight.grad_fn is not None, f{name}.weight grad_fn broken7.3 mini_batch 显存线性度确认无 memory leak用 torch.cuda