在日常开发中，经常需要为 PyTorch 训练脚本快速添加早停、混合精度或分布式支持。手动修改容易遗漏逻辑、弄错上下文、破坏原有结构；CodeGeeX 在 VSCode 中能直接理解你的模型代码语义，给出可直接运行的增量补丁，而非笼统建议，这能显著提升开发效率，减少人工调试成本。

安装并启用 CodeGeeX 插件

打开 VSCode 扩展市场（Ctrl+Shift+X），搜索 CodeGeeX，点击安装，重启 VSCode 即可完成初始化。

安装后右下角状态栏会出现“CodeGeeX Ready”提示；如果未显示，说明插件尚未激活，需手动点击状态栏图标 → 选择“Enable for this workspace”。

还有一点需要特别注意：当前工作区必须已初始化 Git 仓库（存在 .git 文件夹）。CodeGeeX 依赖 Git 历史记录分析代码意图，若缺少 Git 仓库，它无法精准定位训练循环上下文，从而影响功能效果。

让 CodeGeeX 理解你的训练脚本结构

打开 train.py，将光标停留在训练主循环（例如 for epoch in range(num_epochs):）的开头处。

按下 Ctrl+Enter（Windows/Linux）或 Cmd+Enter（macOS）触发 CodeGeeX 智能补全。它会自动扫描周围 200 行代码，提取模型定义、数据加载器、优化器、损失函数等关键组件，完成训练上下文识别。

这一步不可省略：如果光标停在注释行或空行，CodeGeeX 会返回“未检测到训练上下文”。必须精确顶在 for/while 循环第一行的冒号前或缩进起始列，才能确保正确识别。

添加早停机制（Early Stopping）

光标保持在训练循环内，输入自然语言指令：

方法一：直接补全
在循环体开头新起一行，输入 // add early stopping with patience=7 and min_delta=1e-4，按 Ctrl+Enter。
CodeGeeX 会插入完整的早停类实例化、验证损失监控、best_loss 初始化以及 break 退出逻辑。

方法二：交互式编辑
选中已有的验证逻辑块（包含 val_loss 计算），右键 → “CodeGeeX: Refactor with AI” → 选择模板“Add Early Stopping”。
需要注意的是，生成代码默认使用 torch.sa ve(model.state_dict(), ...)，不会保存 optimizer 或 scheduler。如果需要断点续训，必须手动补全保存 optimizer 的逻辑。

启用混合精度训练（AMP）

第一步：在 import 区末尾插入一行：from torch.cuda.amp import autocast, GradScaler。

第二步：在训练循环外初始化 scaler：scaler = GradScaler()。

第三步：把原来的训练步骤包裹进 autocast 上下文，并用 scaler.step() 替代 optimizer.step()：
将 loss.backward() 和 optimizer.step() 之间的代码替换为：
scaler.scale(loss).backward() → scaler.step(optimizer) → scaler.update()。

注意：验证阶段必须显式关闭 autocast（加上 with torch.no_grad():），否则有可能触发 CUDA error: invalid device ordinal 错误。

一键生成分布式训练适配代码

把光标放到 if __name__ == "__main__": 下方，输入指令：
Convert to DDP training using torch.distributed.launch。

CodeGeeX 会自动生成四段代码：① 初始化 init_process_group；② 将 model 包装为 DDP；③ 修改 DataLoader 的 sampler 为 DistributedSampler；④ 添加 torch.distributed.destroy_process_group() 收尾。

生成后立刻检查 os.environ["MASTER_PORT"] 是否被硬编码为 '29500'。如果这个端口被占用，需要手动改成其他未被占用的值，否则多卡启动必定失败。