[DIPU] _amp_update_scale_算子未对dim=0的tensor做判断处理

[Reinerzhou]

背景

export DIPU_MOCK_CUDA=True
在运行llama_finetune时遇到_amp_update_scale_算子会出现报错。

问题描述

在export DIPU_MOCK_CUDA=True的情况下执行以下代码：
`import torch
import torch_dipu

from torch import tensor

_scale = tensor(65536., device='cuda')
found_inf_combined = tensor(0., device='cuda')
_growth_tracker = tensor(0, device='cuda', dtype=torch.int32)

_growth_factor = 2.0
_backoff_factor = 0.5
_growth_interval = 2000

torch.amp_update_scale(_scale, _growth_tracker, found_inf_combined, _growth_factor, _backoff_factor, _growth_interval)`

会出现错误：

初步判断是这里的逻辑没有对dim=0的输入tensor做处理：

https://github.com/DeepLink-org/deeplink.framework/blob/16e155d65f2a5e56d703b3e6acf3d9036b5acb1b/dipu/torch_dipu/csrc_dipu/aten/ops/CustomFallbackFunctionsForAmpGradScaler.cpp#L74C1-L103C2