AmpUpdateScale

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：实现AMP（Automatic Mixed Precision）训练中的动态Scale更新逻辑。根据当前scale值、growth tracker计数器以及是否发现Inf/NaN，动态调整loss scale大小。

• 计算公式：

  $$\text{updated\_scale}=\{\begin{array}{ll}\text{current\_scale}\times \text{backoff\_factor}& \text{if found\_inf}\ne 0\\ \text{current\_scale}\times \text{growth\_factor}& \text{if growth\_tracker + 1 = growth\_interval and new\_scale is finite}\\ \text{current\_scale}& \text{otherwise}\end{array}$$

  $$\text{updated\_growth\_tracker}=\{\begin{array}{ll}0& \text{if found\_inf}\ne 0\text{ or growth triggered}\\ \text{growth\_tracker}+1& \text{otherwise}\end{array}$$

  其中：

  • $\text{current\_scale}$：当前的loss scale值
  • $\text{found\_inf}$：是否检测到Inf/NaN的标志（0表示正常，非0表示发现Inf/NaN）
  • $\text{growth\_tracker}$：连续未出现Inf/NaN的步数计数器
  • $\text{growth\_factor}$：scale增长因子（通常为2.0）
  • $\text{backoff\_factor}$：scale回退因子（通常为0.5）
  • $\text{growth\_interval}$：触发scale增长的间隔步数

• 使用场景：AMP训练中的动态损失缩放（Dynamic Loss Scaling），用于在FP16/BF16混合精度训练中防止梯度下溢。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
current_scale	输入	当前的loss scale值。	FLOAT、FLOAT16、BFLOAT16	ND
growth_tracker	输入	连续未出现Inf/NaN的步数计数器。	INT32	ND
found_inf	输入	是否检测到Inf/NaN的标志，0表示正常，非0表示发现Inf/NaN。	FLOAT、FLOAT16、BFLOAT16	ND
growth_factor	属性	scale增长因子，通常设置为2.0。	FLOAT	-
backoff_factor	属性	scale回退因子，通常设置为0.5。	FLOAT	-
growth_interval	属性	触发scale增长的间隔步数。	INT	-
updated_scale	输出	更新后的loss scale值。	FLOAT、FLOAT16、BFLOAT16	ND
updated_growth_tracker	输出	更新后的growth tracker计数器。	INT32	ND

• shape约束：所有输入输出张量均为标量，shape为 [1]。
• 数据类型约束：current_scale与found_inf、updated_scale的数据类型需一致。

约束说明

无