能，但需先 torch.compile(model)，再在训练 step 中用 torch.cuda.amp.autocast() 包裹前向，gradscaler 必须配合使用且顺序为 scale(loss).backward() → scaler.step(optimizer) → optimizer.zero_grad() → scaler.update()。
PyTorch 2.x 中 torch.compile 和 amp 能否一起用？
能，但顺序和位置很关键。混合精度（AMP）必须放在 torch.compile 封装之后、训练循环之前启用，否则编译器可能无法正确追踪 autocast 和 GradScaler 的行为，导致梯度缩放失效或 RuntimeError。
常见错误现象：RuntimeError: expected scalar type Half but found Float 或 loss 突然变为 NaN —— 多半是 torch.compile(model) 后又在 model 内部手动调用 model.half()，破坏了编译图的类型一致性。
-
torch.compile默认不改变模型 dtype，它只优化计算图；混合精度由torch.cuda.amp.autocast动态控制算子精度，二者职责分离 - 推荐写法：先
model = torch.compile(model)，再在训练 step 中用with torch.cuda.amp.autocast():包裹前向 -
GradScaler必须作用于optimizer.step()前，且 scaler 的scale(loss).backward()要在 autocast 块内或紧随其后（不能跨 compile 边界）
为什么 torch.cuda.amp.GradScaler 在 PyTorch 2.x 里仍不可省略？
即使启用了 torch.compile，GradScaler 依然必要——它负责检测和跳过含 inf/nan 梯度的 step，并动态调整 loss scale。编译器不接管这部分逻辑。
容易踩的坑：scaler.step(optimizer) 返回 None（而非 bool），且不会自动 zero_grad；若忘记手动调用 optimizer.zero_grad()，会导致梯度累积异常。
- 必须配对使用：
scaler.scale(loss).backward()→scaler.step(optimizer)→optimizer.zero_grad() -
scaler.update()必须在scaler.step()之后调用，否则下一轮 scale 不会自适应调整 - 若使用 DDP，
GradScaler无需额外适配，它已内置对 all-reduce 后梯度的 scale 处理
哪些层或操作在 autocast 下默认不降精度？
autocast 对部分算子有硬编码白名单（如 torch.nn.functional.embedding、torch.nn.LayerNorm 输入/输出默认保持 float32），这是为数值稳定性做的妥协，不是 bug。
典型场景：Transformer 类模型中，如果 embedding 输出直接进 softmax，float16 下易出现 softmax overflows（logits 过大 → exp 结果 inf）；autocast 自动保留这部分为 float32。
- 可通过
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True启用更激进的 fp16 reduction（仅限 Ampere+ GPU） - 若想强制某 op 用 float16，可用
torch.set_default_dtype(torch.float16)，但需自行验证数值稳定性 - 自定义 CUDA 扩展或第三方算子（如 flash-attn）需显式声明支持 fp16，否则 autocast 会 fallback 到 float32
如何验证混合精度是否真正生效？
不能只看模型参数 dtype（next(model.parameters()).dtype 仍是 torch.float32），要看实际 kernel 启动的精度。最可靠方式是用 torch.profiler 抓取 CUDA kernel 名称。
实操建议：在训练 step 中插入 profiler，过滤 kernel 名含 half 或 f16 的条目：
with torch.profiler.profile(record_shapes=True, with_flops=True) as prof:
with torch.cuda.amp.autocast():
out = model(x)
prof.key_averages().table(row_limit=10, sort_by="flops")
若看到大量 cublasLtMatmul 或 __nv_cublasLtMatmul（而非 cublasGemmEx），说明已走 fp16 路径。同时注意 aten::softmax 等算子若显示 float 而非 half，属正常设计行为。
一个常被忽略的点：tensor 的 device 必须是 CUDA，autocast 在 CPU 上完全无效；哪怕 model 已 .cuda()，也要确认输入 tensor 也 on cuda，否则 silently fallback 到全 float32。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
