feat: save forward autocast context in FunctionCtx

[Chamberlain0w0]

为 autograd::FunctionCtx 增加了用于保存并恢复前向执行时的 autocast state (放在框架中是 AutocastContext 结构体）的功能，目标是对齐 PyTorch torch.amp.custom_fwd/custom_bwd 的语义：前向阶段记录当前生效的 AutocastContext；当自定义 backward 或重计算路径需要重新执行部分前向子图时，可以显式恢复该前向 AutocastContext。

改动内容

1. Autocast 侧修改：
   • 新增 GetCurrentAutocastContext()，用于快照当前线程本地的 autocast 上下文。
   • 新增 AutocastGuard(const AutocastContext&) 的构造方法，调用方可以通过 RAII 语义临时恢复一个已保存的 autocast 上下文。
2. Function 侧修改：
   • 在执行 Function::Forward() 之前，将前向 autocast 上下文保存到 FunctionCtx。
   • 新增 FunctionCtx::GetForwardAutocastContext()，供自定义 backward 和重计算路径获取前向阶段保存的 autocast 上下文。
   • 在重置 FunctionCtx 保存变量时，同时重置已保存的前向 autocast 上下文。
   • 注意：并不会自动将所有 backward 都包裹在前向 autocast 上下文中。保存的上下文通过 FunctionCtx 暴露，只有显式重算前向计算的代码路径才会通过上述接口来主动恢复该上下文。
3. 新增 autograd 测试，覆盖以下场景：
   • 前向 autocast 上下文会被正确记录到 FunctionCtx；
   • 手动恢复后，能够观察到与前向一致的 enabled/device/dtype；
   • 退出 restore guard 后，调用方原有的 autocast 上下文会被正确恢复。

PyTorch 对齐证据

torch 没有显式包出一个 AutoCastContext 结构体，但是的确存在 forward 会存下此 context 然后 backward 再度调用恢复的情况。对于常见的两种情况：

1. custom backward：相关逻辑实现位于 torch/amp/autocast_mode.py。

   args[0] 就是 FunctionCtx。自定义前向会存下来相应信息，反向再用相应信息开启 autocast。

2. 重计算：相关逻辑实现位于 torch/utils/checkpoint.py。

   有一个存 autocast 信息的 helper，它在 Checkpoint.Function.forward() 会被调用然后记录信息，在 Checkpoint.Function.backward() 重新恢复。

[chen2021673]

LGTM

[kilinchange]

命名按取设值函数来吧，forward_autocast_context，下面 SaveForwardAutocastContext 改成 set_forward_autocast_context