PluginPattern支持nn.Module级别的匹配与替换

[LeoLau94]

Describe the feature🚀

我们希望Plugin Pattern能够实现如下效果，注意，如果需要共享参数，应由用户在构造时完成共享。

                class TestModule(torch.nn.Module):
                    def __init__(self, *args, **kwargs):
                        super().__init__(*args, **kwargs)
                        self.w = torch.nn.Parameter(torch.rand(1024, 1024))
                        self.bias = torch.Tensor(1024)

                    def forward(self, x):
                        return (x @ self.w + self.bias,)

                class AnotherModule(torch.nn.Module):
                    def __init__(self, weight, *args, **kwargs):
                        super().__init__(*args, **kwargs)
                        self.w = torch.nn.Parameter(weight)

                    def forward(self, x):
                        return x + self.w

                weight = torch.rand(1024, 1024)
                register_plugin_pattern(
                    TestModule(), (torch.empty(1024, 1024),), AnotherModule(weight), target=Target.none
                )
                config = XpuGraphConfig(
                    is_training=False,
                    target=Target.none,
                    vendor_compiler_config={},
                    debug=True,
                    enable_cache=False,
                    freeze=True,
                )
                compiler = XpuGraph(config)

                input_tensor = torch.rand(1024, 1024)
                with torch.inference_mode():
                    compiled = torch.compile(TestModule(), fullgraph=True, dynamic=False, backend=compiler)
                    assert is_similar(compiled(input_tensor)[0], AnotherModule(weight)(input_tensor)) is True

在不开启parameter freezing的场景下

class Matmul:
       def __init__(self, w):
              self.w = w
       def forward(self, x):
              return self.w @ x

class Replace:
       def __init__(self, w):
              self.w = w
       def forward(self, x):
              return self.w + x

w = torch.empty(1024,1024)
# 对于Matmul(w) -> Replace(w)
# pattern和replace共享某些参数，我们建议使用如下方式match&rewrite
def matmul(x, w):
     return x@w

def replace(x, w):
     return x+w

# 修改对应的module实现来实现代码复用
class Matmul:
       def __init__(self, w):
              self.w = w
       def forward(self, x):
              return matmul(x,self.w)

class Replace:
       def __init__(self, w):
              self.w = w
       def forward(self, x):
              return replace(x, self.w)

# 即使不共享参数，在不使能paramter freezing下，pattern也应该只注册function 而不是module

def matmul(x, w):
     return x@w

class Matmul:
       def __init__(self, w):
              self.w = w
       def forward(self, x):
              return matmul(x,self.w)

# replacement即可以注册function，也可以注册module

def replace(x, w):
     return x+w

class Replace:
       def __init__(self, w):
              self.w = w
       def forward(self, x):
              return x + self.w

References

https://github.com/pytorch/pytorch/blob/main/torch/_functorch/_aot_autograd/graph_capture_wrappers.py#L1323

Environment details

PyTorch: 2.5.1
torch_npu : 2.5.1
XpuGraph: 0.4.0

[LeoLau94]

我们注意到，默认情况下，dynamo会将本该是get_attr的Op，全部按照placeholder来处理

这是因为dynamo只有在使能parameter freezing的情况下，才会按照get_attr来处理

在xpu_graph中，我们提供了freeze和is_training来控制dynamo使能parameter freezing

这个问题的原因大抵如下：

在fx graph种，get_attr表示从graph自身的Module中去取对应的target，因此，如果使用get_attr，需要将对应的变量浅拷贝至GraphModule中。如下图所示

实际上这大概等价于

graph_module._named_paramter['self__parameter_w`] = nn_module.w

“推理场景”的寓意一般来说是“无需梯度但不禁止用户替换对应的parameter/tensor”，这种情况下参数并不是真正的freezed。

目前存在“无需梯度”但“需要更新权重”的场景，如果dynamo还是使用get_attr的方式去读取对应的参数并且用户没有使用原位操作来更新对应的参数，那么Graph就没法得到最新的变量。

因此dynamo在处理这种情况，为了避免引入额外的用户使用限制（即必须使用原位操作来更新参数），不会使用get_attr，而是将其转换为placeholder，大抵等价于

graph_context.module = nn_module
graph_module.forward(graph_context.module.w)

这种限制还和异步并行相关，因为原位操作一般来说会引入同步操作。

[LeoLau94]

目前开发中PluginPattern只实现了parameter freezing下的module替换，对于其他情况，我们强烈建议用户仍然走functional这种传统的match&rewrite编写方式， 即

(code by @zhangjihang-BD )

import torch
from abc import ABC


# 需要在module以外去定义一个forward函数（不能是成员函数）
def out_forward(x, weight):
    return torch.matmul(x, weight)

torch._dynamo.allow_in_graph(out_forward)

class Foo(ABC, torch.nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self._weight = torch.nn.Parameter(torch.randn(in_c, out_c))

    def forward(self, x):
        # 在这里包回去
        return out_forward(x, self._weight)

[LeoLau94]

某种程度上，我们当然能够在Patter的fx graph里，将get_attr手动替换成placeholder节点，从而充分利用placeholder的通配性质，去匹配paramter non-/freezing两种场景。

可是，我们没法让replacment的get_attr全部换成placeholder，这是因为，即使replacement想要继承pattern中匹配的到的参数，在我们的issue给出的例子中，它一般也是由用户传入对应的参数并使用__init__构造好replacement module，再转交给我们。因此，在ewrite时，会出现pattern和replacement没法一一对应的情况。

如果用户不能提前构造好replacement module，我们建议如上文那样，通过funtional方式去做match&rewrite，把module替换视为函数替换。

[LeoLau94]

核心要点代码化为：

weight = torch.Tensor()
module.weight = weight
graph_module.weight = module.weight
weight_1 = torch.Tensor()
module.weight = weight_1
assert graph_module.weight != module.weight
assert graph_module.weight == weight
# If we have & in python, then
weight = torch.Tensor()
module.weight = weight
graph_module.weight = &(module.weight)
weight_1 = torch.Tensor()
module.weight = weight_1
assert graph_module.weight == module.weight
assert graph_module.weight != weight
# That's why dynamo make get_attr as placeholder by default even it's in inference mode.

当然，如果graph module直接把原来的module作为submodule保存起来，也不会有这个问题，但是默认情况下，module在dynamo tracing的时候是会直接展开以获得更大的优化空间。

[LeoLau94]

目前的需求是，替换pattern的replacement需要完成以下事情

1. 动态构造module对象，因为有一些信息来自于匹配的段落；
2. 调用module对象执行前推逻辑；

因此，我们的pattern匹配还需要抓取到构造module对象所需要的参数（作为placeholder表现在图里）

此外，还要避免重复构造module对象

[zhangjihang-BD]

我们注意到，默认情况下，dynamo会将本该是get_attr的Op，全部按照placeholder来处理

这是因为dynamo只有在使能parameter freezing的情况下，才会按照get_attr来处理

在xpu_graph中，我们提供了freeze和is_training来控制dynamo使能parameter freezing

这个问题的原因大抵如下：

在fx graph种，get_attr表示从graph自身的Module中去取对应的target，因此，如果使用get_attr，需要将对应的变量浅拷贝至GraphModule中。如下图所示

实际上这大概等价于

graph_module._named_paramter['self__parameter_w`] = nn_module.w

“推理场景”的寓意一般来说是“无需梯度但不禁止用户替换对应的parameter/tensor”，这种情况下参数并不是真正的freezed。

目前存在“无需梯度”但“需要更新权重”的场景，如果dynamo还是使用get_attr的方式去读取对应的参数并且用户没有使用原位操作来更新对应的参数，那么Graph就没法得到最新的变量。

因此dynamo在处理这种情况，为了避免引入额外的用户使用限制（即必须使用原位操作来更新参数），不会使用get_attr，而是将其转换为placeholder，大抵等价于

graph_context.module = nn_module
graph_module.forward(graph_context.module.w)

这种限制还和异步并行相关，因为原位操作一般来说会引入同步操作。

话说无需梯度，但需要更新权重的场景是啥呢

[LeoLau94]

我们注意到，默认情况下，dynamo会将本该是get_attr的Op，全部按照placeholder来处理

这是因为dynamo只有在使能parameter freezing的情况下，才会按照get_attr来处理
在xpu_graph中，我们提供了freeze和is_training来控制dynamo使能parameter freezing
这个问题的原因大抵如下：
在fx graph种，get_attr表示从graph自身的Module中去取对应的target，因此，如果使用get_attr，需要将对应的变量浅拷贝至GraphModule中。如下图所示
实际上这大概等价于

graph_module._named_paramter['self__parameter_w`] = nn_module.w

“推理场景”的寓意一般来说是“无需梯度但不禁止用户替换对应的parameter/tensor”，这种情况下参数并不是真正的freezed。
目前存在“无需梯度”但“需要更新权重”的场景，如果dynamo还是使用get_attr的方式去读取对应的参数并且用户没有使用原位操作来更新对应的参数，那么Graph就没法得到最新的变量。
因此dynamo在处理这种情况，为了避免引入额外的用户使用限制（即必须使用原位操作来更新参数），不会使用get_attr，而是将其转换为placeholder，大抵等价于

graph_context.module = nn_module
graph_module.forward(graph_context.module.w)

这种限制还和异步并行相关，因为原位操作一般来说会引入同步操作。

话说无需梯度，但需要更新权重的场景是啥呢

有些场景可能会基于新数据进行实时训练（微调），然后部署替换掉旧的权重