| 1 | import torch |
| 2 | |
| 3 | from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict |
| 4 | from ltx_core.quantization.fp8_cast import calculate_weight_float8 |
| 5 | from ltx_core.quantization.fp8_scaled_mm import quantize_weight_to_fp8_per_tensor |
| 6 | |
| 7 | |
| 8 | def apply_loras( |
| 9 | model_sd: StateDict, |
| 10 | lora_sd_and_strengths: list[LoraStateDictWithStrength], |
| 11 | dtype: torch.dtype | None = None, |
| 12 | destination_sd: StateDict | None = None, |
| 13 | ) -> StateDict: |
| 14 | sd = {} |
| 15 | if destination_sd is not None: |
| 16 | sd = destination_sd.sd |
| 17 | size = 0 |
| 18 | device = torch.device("meta") |
| 19 | inner_dtypes = set() |
| 20 | for key, weight in model_sd.sd.items(): |
| 21 | if weight is None: |
| 22 | continue |
| 23 | # Skip scale keys - they are handled together with their weight keys |
| 24 | if key.endswith(".weight_scale"): |
| 25 | continue |
| 26 | device = weight.device |
| 27 | target_dtype = dtype if dtype is not None else weight.dtype |
| 28 | deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16 |
| 29 | |
| 30 | scale_key = key.replace(".weight", ".weight_scale") if key.endswith(".weight") else None |
| 31 | is_scaled_fp8 = scale_key is not None and scale_key in model_sd.sd |
| 32 | |
| 33 | deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, device) |
| 34 | fused = _fuse_deltas(deltas, weight, key, sd, target_dtype, device, is_scaled_fp8, scale_key, model_sd) |
| 35 | |
| 36 | sd.update(fused) |
| 37 | for tensor in fused.values(): |
| 38 | inner_dtypes.add(tensor.dtype) |
| 39 | size += tensor.nbytes |
| 40 | |
| 41 | if destination_sd is not None: |
| 42 | return destination_sd |
| 43 | return StateDict(sd, device, size, inner_dtypes) |
| 44 | |
| 45 | |
| 46 | def _prepare_deltas( |
| 47 | lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device |
| 48 | ) -> torch.Tensor | None: |
| 49 | deltas = [] |
| 50 | prefix = key[: -len(".weight")] |
| 51 | key_a = f"{prefix}.lora_A.weight" |
| 52 | key_b = f"{prefix}.lora_B.weight" |
| 53 | for lsd, coef in lora_sd_and_strengths: |
| 54 | if key_a not in lsd.sd or key_b not in lsd.sd: |
| 55 | continue |
| 56 | a = lsd.sd[key_a].to(device=device) |
| 57 | b = lsd.sd[key_b].to(device=device) |
| 58 | product = torch.matmul(b * coef, a) |
| 59 | del a, b |
| 60 | deltas.append(product.to(dtype=dtype)) |
| 61 | if len(deltas) == 0: |
| 62 | return None |
| 63 | elif len(deltas) == 1: |
| 64 | return deltas[0] |
| 65 | return torch.sum(torch.stack(deltas, dim=0), dim=0) |
| 66 | |
| 67 | |
| 68 | def _fuse_deltas( |
| 69 | deltas: torch.Tensor | None, |
| 70 | weight: torch.Tensor, |
| 71 | key: str, |
| 72 | sd: dict[str, torch.Tensor], |
| 73 | target_dtype: torch.dtype, |
| 74 | device: torch.device, |
| 75 | is_scaled_fp8: bool, |
| 76 | scale_key: str | None, |
| 77 | model_sd: StateDict, |
| 78 | ) -> dict[str, torch.Tensor]: |
| 79 | if deltas is None: |
| 80 | if key in sd: |
| 81 | return {} |
| 82 | fused = _copy_weight_without_lora(weight, key, target_dtype, device, is_scaled_fp8, scale_key, model_sd) |
| 83 | elif weight.dtype == torch.float8_e4m3fn: |
| 84 | if is_scaled_fp8: |
| 85 | fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd) |
| 86 | else: |
| 87 | fused = _fuse_delta_with_cast_fp8(deltas, weight, key, target_dtype, device) |
| 88 | elif weight.dtype == torch.bfloat16: |
| 89 | fused = _fuse_delta_with_bfloat16(deltas, weight, key, target_dtype) |
| 90 | else: |
| 91 | raise ValueError(f"Unsupported dtype: {weight.dtype}") |
| 92 | |
| 93 | return fused |
| 94 | |
| 95 | |
| 96 | def _copy_weight_without_lora( |
| 97 | weight: torch.Tensor, |
| 98 | key: str, |
| 99 | target_dtype: torch.dtype, |
| 100 | device: torch.device, |
| 101 | is_scaled_fp8: bool, |
| 102 | scale_key: str | None, |
| 103 | model_sd: StateDict, |
| 104 | ) -> dict[str, torch.Tensor]: |
| 105 | """Copy original weight (and scale if applicable) when no LoRA affects this key.""" |
| 106 | result = {key: weight.clone().to(dtype=target_dtype, device=device)} |
| 107 | if is_scaled_fp8: |
| 108 | result[scale_key] = model_sd.sd[scale_key].clone() |
| 109 | return result |
| 110 | |
| 111 | |
| 112 | def _fuse_delta_with_scaled_fp8( |
| 113 | deltas: torch.Tensor, |
| 114 | weight: torch.Tensor, |
| 115 | key: str, |
| 116 | scale_key: str, |
| 117 | model_sd: StateDict, |
| 118 | ) -> dict[str, torch.Tensor]: |
| 119 | """Dequantize scaled FP8 weight, add LoRA delta, and re-quantize.""" |
| 120 | weight_scale = model_sd.sd[scale_key] |
| 121 | |
| 122 | original_weight = weight.t().to(torch.float32) * weight_scale |
| 123 | |
| 124 | new_weight = original_weight + deltas.to(torch.float32) |
| 125 | |
| 126 | new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight) |
| 127 | return {key: new_fp8_weight, scale_key: new_weight_scale} |
| 128 | |
| 129 | |
| 130 | def _fuse_delta_with_cast_fp8( |
| 131 | deltas: torch.Tensor, |
| 132 | weight: torch.Tensor, |
| 133 | key: str, |
| 134 | target_dtype: torch.dtype, |
| 135 | device: torch.device, |
| 136 | ) -> dict[str, torch.Tensor]: |
| 137 | """Fuse LoRA delta with cast-only FP8 weight (no scale factor).""" |
| 138 | if str(device).startswith("cuda"): |
| 139 | deltas = calculate_weight_float8(deltas, weight) |
| 140 | else: |
| 141 | deltas.add_(weight.to(dtype=deltas.dtype, device=device)) |
| 142 | return {key: deltas.to(dtype=target_dtype)} |
| 143 | |
| 144 | |
| 145 | def _fuse_delta_with_bfloat16( |
| 146 | deltas: torch.Tensor, |
| 147 | weight: torch.Tensor, |
| 148 | key: str, |
| 149 | target_dtype: torch.dtype, |
| 150 | ) -> dict[str, torch.Tensor]: |
| 151 | """Fuse LoRA delta with bfloat16 weight.""" |
| 152 | deltas.add_(weight) |
| 153 | return {key: deltas.to(dtype=target_dtype)} |
| 154 |