Spaces:
Paused
Paused
File size: 5,354 Bytes
8cf98bd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
"""
References:
https://github.com/KONAKONA666/q8_kernels/blob/9cee3f3d4ca5ec8ab463179be32c8001e31f8f33/q8_kernels/utils/convert_weights.py
"""
import torch
import torch.nn as nn
from q8_kernels.modules.rms_norm import RMSNorm as QRMSNorm
from diffusers.models.normalization import RMSNorm
from q8_kernels.modules.activations import GELU as QGELU
from diffusers.models.activations import GELU
from q8_kernels.modules.linear import Q8Linear
from q8_attention_processors import LTXVideoQ8AttentionProcessor
MODULES_TO_NOT_CONVERT = ["proj_in", "time_embed", "caption_projection", "proj_out"]
def replace_linear(model, current_key_name=None, replaced=False):
for name, child in model.named_children():
if current_key_name is None:
current_key_name = []
current_key_name.append(name)
if isinstance(child, nn.Linear) and name not in MODULES_TO_NOT_CONVERT:
# Check if the current key is not in the `modules_to_not_convert`
current_key_name_str = ".".join(current_key_name)
if not any(
(key + "." in current_key_name_str) or (key == current_key_name_str) for key in MODULES_TO_NOT_CONVERT
):
new_linear = Q8Linear(
child.in_features, child.out_features, bias=child.bias is not None, device=child.weight.device
)
setattr(model, name, new_linear)
replaced = True
else:
replace_linear(model=child, current_key_name=current_key_name, replaced=replaced)
current_key_name.pop(-1)
return model, replaced
def get_parent_module_and_attr(root, dotted_name: str):
"""
Splits 'a.b.c' into:
- parent module = root.a.b
- attr_name = 'c'
"""
parts = dotted_name.split(".")
*parent_parts, attr_name = parts
parent_module = root
for p in parent_parts:
parent_module = getattr(parent_module, p)
return parent_module, attr_name
def replace_rms_norm(model):
modules_to_replace = []
for dotted_name, module in model.named_modules():
if isinstance(module, RMSNorm):
modules_to_replace.append((dotted_name, module))
replaced = False
for dotted_name, module in modules_to_replace:
parent, attr_name = get_parent_module_and_attr(model, dotted_name)
new_norm = QRMSNorm(
dim=module.dim,
elementwise_affine=module.elementwise_affine,
)
setattr(parent, attr_name, new_norm)
replaced = True
return model, replaced
def replace_gelu(model, replaced=False):
for name, child in model.named_children():
if isinstance(child, GELU):
new_gelu = QGELU(
dim_in=child.proj.in_features,
dim_out=child.proj.out_features,
approximate=child.approximate,
bias=child.proj.bias is not None,
)
setattr(model, name, new_gelu)
replaced = True
else:
replace_gelu(model=child, replaced=replaced)
return model, replaced
def set_attn_processors(model, processor):
def fn_recursive_attn_processor(name, module: torch.nn.Module, processor):
if hasattr(module, "set_processor"):
module.set_processor(processor)
for sub_name, child in module.named_children():
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
for name, module in model.named_children():
fn_recursive_attn_processor(name, module, processor)
def attn_processors(model) -> dict:
processors = {}
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict):
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
for sub_name, child in module.named_children():
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
return processors
for name, module in model.named_children():
fn_recursive_add_processors(name, module, processors)
return processors
def check_transformer_replaced_correctly(model):
for block in model.transformer_blocks:
assert isinstance(block.attn1.to_q, Q8Linear), f"{type(block.attn1.to_q)=} not linear."
assert isinstance(block.attn2.to_q, Q8Linear), f"{type(block.attn2.to_q)=} not linear."
assert block.attn1.to_q.weight.dtype == torch.int8, f"{block.attn1.to_q.weight.dtype=}."
assert block.attn2.to_q.weight.dtype == torch.int8, f"{name=} {block.attn2.to_q.weight.dtype=}."
for name, module in model.named_modules():
if "norm" in name and "norm_out" not in name:
assert isinstance(module, QRMSNorm), f"{name=}, {type(module)=}"
for block in model.transformer_blocks:
assert isinstance(block.ff.net[0], QGELU), f"{type(block.ff.net[0])=}"
if getattr(block.ff.net[0], "proj", None) is not None:
assert block.ff.net[0].proj.weight.dtype == torch.int8, f"{block.ff.net[0].proj.weight.dtype=}."
set_attn_processors(model, LTXVideoQ8AttentionProcessor())
all_attn_processors = attn_processors(model)
for k, v in all_attn_processors.items():
assert isinstance(v, LTXVideoQ8AttentionProcessor), f"{name} is not of type LTXVideoQ8AttentionProcessor."
|