Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add context manager TensorInplaceAssign #396

Merged
merged 33 commits into from
Dec 13, 2023

Conversation

marigoold
Copy link
Contributor

@marigoold marigoold commented Dec 4, 2023

在 diffusers 中,load LoRA 的实现是直接把 fuse 之后的 weight 赋值给 self.weight(比如这里),这会使赋值后 tensor 的内存地址改变,而 graph 依然使用原来的内存地址,导致 LoRA 无法正确加载。
这里实现了一个 Tensor 和 Parameter 的子类,重写了 property data 的 getter 和 setter,使其 set 的时候原地拷贝,不改变内存地址。

正确性验证:

  1. 在每个赋值语句后面都加了内存地址不变的断言,没有挂掉
        dptr = self.regular_linear_layer.weight.data.data_ptr()
        self.regular_linear_layer.weight.data = fused_weight.to(device=device, dtype=dtype)
        if "AutoInplaceCopy" in str(type(self.regular_linear_layer.weight.data)):
            print(f"got type [{type(self.regular_linear_layer.weight.data)}], shape is {self.regular_linear_layer.weight.data.shape}")
            assert self.regular_linear_layer.weight.data.data_ptr() == dptr
  1. examples/text_to_image_sdxl_lora.py 脚本中,先用 oneflow graph forward 一次,再加载 LoRA,结果与 torch 相同;在 oneflow_compile 中删掉对应转换语句,结果与 torch 不同。

@marigoold marigoold changed the title [WIP] add AutoInplaceCopyTensor add context manager AutoInplaceAssign Dec 5, 2023
dptr1 = eager.linear1.weight.data.data_ptr()
dptr2 = eager.linear2.weight.data.data_ptr()

with AutoInplaceAssign(eager):
Copy link
Collaborator

@strint strint Dec 5, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

你用这个解决下 lora un-fuse 那个问题看看

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AutoInplaceAssign -> TensorInplaceAssign

Auto 没有什么含义

@marigoold marigoold changed the title add context manager AutoInplaceAssign add context manager TensorInplaceAssign Dec 7, 2023
@marigoold marigoold marked this pull request as ready for review December 8, 2023 04:03
@strint
Copy link
Collaborator

strint commented Dec 8, 2023

这个 PR 需要配合一个 lora unfuse 的测试吧

@strint
Copy link
Collaborator

strint commented Dec 8, 2023

#407 合并到这里,并增加下那个 unfuse lora 的测试吧

@marigoold marigoold merged commit 4123a7a into main Dec 13, 2023
4 checks passed
@marigoold marigoold deleted the dev_wy_add_AutoInplaceCopyTensor branch December 13, 2023 07:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants