nodes/kijai-ComfyUI-PyramidFlowWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

allow torch.compiling more stuff

Browse Source
This commit is contained in:
kijai
2024-11-13 23:59:10 +02:00
parent acb19bcd29
commit 723537bb75
4 changed files with 26 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fp8_optimization.py
+6 -5
View File
@@ -36,10 +36,11 @@ def fp8_linear_forward(cls, original_dtype, input):
else:
return cls.original_forward(input)
def convert_fp8_linear(module, original_dtype):
def convert_fp8_linear(module, original_dtype, params_to_keep):
setattr(module, "fp8_matmul_enabled", True)
for name, module in module.named_modules():
if isinstance(module, nn.Linear):
original_forward = module.forward
setattr(module, "original_forward", original_forward)
setattr(module, "forward", lambda input, m=module: fp8_linear_forward(m, original_dtype, input))
if not any(keyword in name for keyword in params_to_keep):
if isinstance(module, nn.Linear):
original_forward = module.forward
setattr(module, "original_forward", original_forward)
setattr(module, "forward", lambda input, m=module: fp8_linear_forward(m, original_dtype, input))