[tests] feat: add AoT compilation tests #12203
Conversation
|
|
||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| package_path = os.path.join(tmpdir, f"{self.model_class.__name__}.pt2") | ||
| _ = torch._inductor.aoti_compile_and_package(exported_model, package_path=package_path) |
There was a problem hiding this comment.
Not passing in a path should also automatically give you a path in the tmp dir!
|
|
||
| with torch.no_grad(): | ||
| _ = model(**inputs_dict) | ||
| _ = model(**inputs_dict) |
There was a problem hiding this comment.
Is there a particular reason why you're running it twice?
There was a problem hiding this comment.
To emulate the real scenario as the model is typically invoked more than once during the actual generation process.
|
Very nice! Ideally we should pass multiple inputs to the exported (and/or compiled) model in order to test proper dynamism (for instance in the case of |
We already have this kind of test: diffusers/tests/models/test_modeling_common.py Line 2119 in 91a151b I think it's fine for now to just test for static shapes. If dynamic AoT becomes more popular, we can revisit it. |
Very nice! I think that |
What does this PR do?
AoT compilation is exciting because it helps cut the framework overhead. It also helps realize similar benefits as JiT compilation for environments where JiT might not be feasible (ZeroGPU Spaces, for example). For example, with AoT-compilation, we were able to obtain 1.7x speedups in latency on a ZeroGPU Space. Wouldn't have been feasible, otherwise.
To run,
RUN_SLOW=1 RUN_COMPILE=yes pytest tests/models/ -k "aot".