search

facebookincubator / AITemplate

AITemplate is a Python framework which renders neural network into high performance CUDA/HIP C++ code. Specialized for FP16 TensorCore (NVIDIA GPU) and MatrixCore (AMD GPU) inference.
Apache License 2.0
4.56k stars 370 forks source link

[ask] is it possible to copy weights from PyTorch model to complied AIT model #407

Open YaoAstro opened 1 year ago

YaoAstro commented 1 year ago

I have a group of fine-tuned stable diffusion models using the same training process. After compiling one model into AIT format, is it possible to reuse the "selected" AIT graph for other models by directly copying weights over (as opposed to compiling all the other models)?

khabinov commented 1 year ago

Hi @YaoAstro,

Such functionality should be more or less working though not very well documented. However, you can refer to this test to get some idea https://github.com/facebookincubator/AITemplate/blob/32a846507b129003c26600b1bcbe98f93aebcefe/tests/unittest/compiler/test_constant_folding.py#L543.

cc @muchulee8

YaoAstro commented 1 year ago

Thanks for the pointer, @khabinov!