Dtype gen pass

[WHoutstanding]

PR Category

other

Description

对以下算子进行特殊处理
torch.matmul,
torch.nn.functional.linear,
torch.nn.functional.conv2d,
torch.bmm,
torch.nn.functional.scaled_dot_product_attention,

[paddle-bot]

Thanks for your contribution!

[WHoutstanding]

float16和bfloat16都失败的samples原来共有19个，特殊处理后成功转换11个samples
另外8个失败的samples不会在以下算子出错

torch.matmul,
torch.nn.functional.linear,
torch.nn.functional.conv2d,
torch.bmm,
torch.nn.functional.scaled_dot_product_attention,

但是会有新的node类型为call_function的算子或者node类型为call_method的节点Error

[WHoutstanding]

测试samples：

samples/transformers-auto-model/hf-tiny-model-private_tiny-random-Swinv2ForImageClassification
samples/torchgeometric/RECT_L
samples/transformers-auto-model/google_byt5_base
samples/transformers-auto-model/Fsoft-AIC_videberta-base
samples/mmseg/SegNeXt-b
samples/transformers-auto-model/OFA-Sys_chinese-clip-vit-large-patch14
samples/timm/vit_small_patch16_rope_mixed_ape_224.naver_in1k
samples/transformers-auto-model/google/t5-efficient-large-kv128
samples/transformers-auto-model/MoritzLaurer_xtremedistil-l6-h256-mnli-fever-anli-ling-binary
samples/mmseg/CCNet_R101
samples/transformers-auto-model/facebook_sam-vit-large
samples/transformers-auto-model/ogoshi2000_stance-nystromformer
samples/transformers-auto-model/TinyLlama/TinyLlama-1.1B-Chat-v0.4
samples/transformers-auto-model/google-t5_t5-large
samples/transformers-auto-model/apple_aimv2-huge-patch14-224
samples/transformers-auto-model/microsoft_swin-base-patch4-window12-384-in22k
samples/mmpose/pose_swin_b
samples/transformers-auto-model/all-mpnet-base-v2
samples/transformers-auto-model/Neurora_opus-tatoeba-heb-eng

运行log：
call_function_f16_bf16_error_log_init.txt

[Xreki]

日志中，大部分样本报错跟matmul_3 = attn_9 @ v_1、denominator = coef_1.bmm(bmm_2)有关，其中@实际应该是matmul算子，这两个可以再看下。

你这个PR处理的是amp的白名单算子，其他的报错跟amp的黑名单算子有关，可以先不继续。

[WHoutstanding]

手动设置如下白名单：
AMP_CALL_FUNCTION = {
torch.matmul,
torch.mm,
torch.bmm,
torch.nn.functional.linear,
torch.nn.functional.conv1d,
torch.nn.functional.conv2d,
torch.nn.functional.conv3d,
torch.nn.functional.scaled_dot_product_attention,
}

AMP_CALL_METHOD = {
"matmul",
"mm",
"bmm",
}

float16和bfloat16都失败的samples原来共有19个，设置白名单后成功转换12个samples

2个samples不需要rewrite：
model_path='samples/transformers-auto-model/Neurora_opus-tatoeba-heb-eng'
model_path='samples/transformers-auto-model/MoritzLaurer_xtremedistil-l6-h256-mnli-fever-anli-ling-binary'

5个samples没有成功转换，error算子为：

    attention_scores_3 = attention_scores_2.masked_fill(invert, -3.4028234663852886e+38);  attention_scores_2 = invert = None
RuntimeError: value cannot be converted to type at::Half without overflow

    freqs_x = unsqueeze @ unsqueeze_1;  unsqueeze = unsqueeze_1 = None
RuntimeError: expected scalar type Float but found Half

    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
RuntimeError: expected scalar type Float but found Half

    coordinates_3 = coordinates_2 @ to_1;  coordinates_2 = to_1 = None
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::Half

    matmul_3 = attn_9 @ v_1;  attn_9 = v_1 = None
RuntimeError: expected scalar type BFloat16 but found Float

matmul_3 = attn_9 @ v_1 算子报错的sample只有1个 :
samples/mmpose/pose_swin_b

[Xreki]

#593 (comment) 中只是以matmul_3 = attn_9 @ v_1为例，表示@这一类，你贴出的freqs_x = unsqueeze @ unsqueeze_1、coordinates_3 = coordinates_2 @ to_1都属于这一类。这类算子没有转换成功的原因是什么呢？

[WHoutstanding]

重新设置白名单：
AMP_CALL_FUNCTION = {
torch.matmul,
torch.mm,
torch.bmm,
operator.matmul,
torch.nn.functional.linear,
torch.nn.functional.conv1d,
torch.nn.functional.conv2d,
torch.nn.functional.conv3d,
torch.nn.functional.scaled_dot_product_attention,
torch.addmm,
torch.einsum,
}

AMP_CALL_METHOD = {
"matmul",
"mm",
"bmm",
}

float16和bfloat16都失败的samples原来共有19个，重新设置白名单后成功转换15个samples

2个samples不需要rewrite：
model_path='samples/transformers-auto-model/Neurora_opus-tatoeba-heb-eng'
model_path='samples/transformers-auto-model/MoritzLaurer_xtremedistil-l6-h256-mnli-fever-anli-ling-binary'

2个samples没有成功转换，失败的为黑名单算子：

    attention_scores_3 = attention_scores_2.masked_fill(invert, -3.4028234663852886e+38);  attention_scores_2 = invert = None
RuntimeError: value cannot be converted to type at::Half without overflow

    return torch.layer_norm(
RuntimeError: expected scalar type Float but found BFloat16

[Xreki]

LGTM

[Xreki]

这个整体方案，就基本是在fx.Graph层面来实现AMP机制了。给你看下关于这个题目，大模型给的答案。

下面给你一个基于 torch.fx.Graph 实现 AMP（Automatic Mixed Precision）机制的可落地方案，偏向编译 / 图级改写思路，而不是运行时 autocast。结合你之前关注的 FX / Dynamo / export / backend 定制，这里重点放在 FX Graph Rewrite AMP。

一、设计目标（图级 AMP）

我们希望在 FX Graph 层面做到：

1. 对算子进行精度分类

• FP16 / BF16 安全算子（如 matmul / conv / linear）
• FP32 保留算子（如 softmax、layernorm、exp、log、reduce）

2. 自动插入 cast

• 输入 cast → 低精度算子
• 关键算子前后 cast 回 FP32

3. 可控、可扩展

• 不修改原模型代码
• 不依赖 torch.cuda.amp.autocast

二、整体架构

nn.Module
   ↓ symbolic_trace / torch.export
fx.GraphModule
   ↓ AMP Graph Pass
重写后的 fx.Graph
   ↓ torch.compile / backend

三、算子精度策略定义

1️⃣ 定义白名单 / 黑名单

FP16_SAFE_FUNCS = {
    torch.matmul,
    torch.mm,
    torch.bmm,
    torch.nn.functional.linear,
    torch.nn.functional.conv2d,
}

FP32_ONLY_FUNCS = {
    torch.nn.functional.softmax,
    torch.nn.functional.layer_norm,
    torch.exp,
    torch.log,
}

实际工程里建议：

• 用 op schema / aten op 名称
• 支持 per-op override

四、核心：FX Graph 重写 AMP

1️⃣ 基础思路

• 对 Graph 做一次 pass：
• 遍历 node
• 判断 node 是否是 FP16 safe
• 插入 to(dtype) cast 节点

2️⃣ AMP Graph Pass 示例

import torch
import torch.fx as fx

class AMPGraphPass:
    def __init__(self, target_dtype=torch.float16):
        self.target_dtype = target_dtype

    def run(self, gm: fx.GraphModule):
        graph = gm.graph

        for node in list(graph.nodes):
            if node.op == "call_function":
                if node.target in FP16_SAFE_FUNCS:
                    self._cast_inputs(graph, node, self.target_dtype)

                elif node.target in FP32_ONLY_FUNCS:
                    self._cast_inputs(graph, node, torch.float32)

        graph.lint()
        gm.recompile()
        return gm

3️⃣ 插入 cast 节点

def _cast_inputs(self, graph, node, dtype):
    with graph.inserting_before(node):
        new_args = []
        for arg in node.args:
            if isinstance(arg, fx.Node):
                cast = graph.call_method(
                    "to",
                    args=(arg,),
                    kwargs={"dtype": dtype}
                )
                new_args.append(cast)
            else:
                new_args.append(arg)

        node.args = tuple(new_args)

[Xreki]

torch当前版本的黑白名单定义在：https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/autocast_mode.cpp

[lixinqi]

放到新的create_call_func_args函数里。
让语义紧凑

[WHoutstanding]

create_call_function和create_call_method都有创建args的过程，因此增加函数create_new_args