| 1 | import torch |
| 2 | |
| 3 | from ltx_core.loader.module_ops import ModuleOps |
| 4 | from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps |
| 5 | from ltx_core.model.transformer.model import LTXModel |
| 6 | |
| 7 | BLOCK_SIZE = 1024 |
| 8 | |
| 9 | |
| 10 | def calculate_weight_float8(target_weights: torch.Tensor, original_weights: torch.Tensor) -> torch.Tensor: |
| 11 | result = _fused_add_round_launch(target_weights, original_weights, seed=0).to(target_weights.dtype) |
| 12 | target_weights.copy_(result, non_blocking=True) |
| 13 | return target_weights |
| 14 | |
| 15 | |
| 16 | def _fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor: |
| 17 | # Lazy import triton - only available on CUDA platforms |
| 18 | import triton # noqa: PLC0415 |
| 19 | |
| 20 | from ltx_core.loader.kernels import fused_add_round_kernel # noqa: PLC0415 |
| 21 | |
| 22 | if original_weight.dtype == torch.float8_e4m3fn: |
| 23 | exponent_bits, mantissa_bits, exponent_bias = 4, 3, 7 |
| 24 | elif original_weight.dtype == torch.float8_e5m2: |
| 25 | exponent_bits, mantissa_bits, exponent_bias = 5, 2, 15 # noqa: F841 |
| 26 | else: |
| 27 | raise ValueError("Unsupported dtype") |
| 28 | |
| 29 | if target_weight.dtype != torch.bfloat16: |
| 30 | raise ValueError("target_weight dtype must be bfloat16") |
| 31 | |
| 32 | # Calculate grid and block sizes |
| 33 | n_elements = original_weight.numel() |
| 34 | grid = (triton.cdiv(n_elements, BLOCK_SIZE),) |
| 35 | |
| 36 | # Launch kernel |
| 37 | fused_add_round_kernel[grid]( |
| 38 | original_weight, |
| 39 | target_weight, |
| 40 | seed, |
| 41 | n_elements, |
| 42 | exponent_bias, |
| 43 | mantissa_bits, |
| 44 | BLOCK_SIZE, |
| 45 | ) |
| 46 | return target_weight |
| 47 | |
| 48 | |
| 49 | def _naive_weight_or_bias_downcast(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]: |
| 50 | """ |
| 51 | Downcast the weight or bias to the float8_e4m3fn dtype. |
| 52 | """ |
| 53 | return [KeyValueOperationResult(key, value.to(dtype=torch.float8_e4m3fn))] |
| 54 | |
| 55 | |
| 56 | def _upcast_and_round( |
| 57 | weight: torch.Tensor, dtype: torch.dtype, with_stochastic_rounding: bool = False, seed: int = 0 |
| 58 | ) -> torch.Tensor: |
| 59 | """ |
| 60 | Upcast the weight to the given dtype and optionally apply stochastic rounding. |
| 61 | Input weight needs to have float8_e4m3fn or float8_e5m2 dtype. |
| 62 | """ |
| 63 | if not with_stochastic_rounding: |
| 64 | return weight.to(dtype) |
| 65 | return _fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed) |
| 66 | |
| 67 | |
| 68 | def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None: |
| 69 | """ |
| 70 | Replace linear.forward and rms_norm.forward with a version that: |
| 71 | - upcasts weight and bias to input's dtype |
| 72 | - returns F.linear or F.rms_norm calculated in that dtype |
| 73 | """ |
| 74 | |
| 75 | layer.original_forward = layer.forward |
| 76 | |
| 77 | def new_linear_forward(*args, **_kwargs) -> torch.Tensor: |
| 78 | # assume first arg is the input tensor |
| 79 | x = args[0] |
| 80 | w_up = _upcast_and_round(layer.weight, x.dtype, with_stochastic_rounding, seed) |
| 81 | b_up = None |
| 82 | |
| 83 | if layer.bias is not None: |
| 84 | b_up = _upcast_and_round(layer.bias, x.dtype, with_stochastic_rounding, seed) |
| 85 | |
| 86 | return torch.nn.functional.linear(x, w_up, b_up) |
| 87 | |
| 88 | layer.forward = new_linear_forward |
| 89 | |
| 90 | |
| 91 | def _amend_forward_with_upcast( |
| 92 | model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0 |
| 93 | ) -> torch.nn.Module: |
| 94 | """ |
| 95 | Replace the forward method of the model's Linear and RMSNorm layers to forward |
| 96 | with upcast and optional stochastic rounding. |
| 97 | """ |
| 98 | for m in model.modules(): |
| 99 | if isinstance(m, (torch.nn.Linear)): |
| 100 | _replace_fwd_with_upcast(m, with_stochastic_rounding, seed) |
| 101 | return model |
| 102 | |
| 103 | |
| 104 | TRANSFORMER_LINEAR_DOWNCAST_MAP = ( |
| 105 | SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP") |
| 106 | .with_kv_operation( |
| 107 | key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast |
| 108 | ) |
| 109 | .with_kv_operation( |
| 110 | key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast |
| 111 | ) |
| 112 | .with_kv_operation( |
| 113 | key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast |
| 114 | ) |
| 115 | .with_kv_operation( |
| 116 | key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast |
| 117 | ) |
| 118 | .with_kv_operation( |
| 119 | key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast |
| 120 | ) |
| 121 | .with_kv_operation( |
| 122 | key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast |
| 123 | ) |
| 124 | .with_kv_operation( |
| 125 | key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast |
| 126 | ) |
| 127 | .with_kv_operation( |
| 128 | key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast |
| 129 | ) |
| 130 | .with_kv_operation( |
| 131 | key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast |
| 132 | ) |
| 133 | .with_kv_operation( |
| 134 | key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast |
| 135 | ) |
| 136 | .with_kv_operation( |
| 137 | key_prefix="transformer_blocks.", key_suffix="ff.net.2.weight", operation=_naive_weight_or_bias_downcast |
| 138 | ) |
| 139 | .with_kv_operation( |
| 140 | key_prefix="transformer_blocks.", key_suffix="ff.net.2.bias", operation=_naive_weight_or_bias_downcast |
| 141 | ) |
| 142 | ) |
| 143 | |
| 144 | UPCAST_DURING_INFERENCE = ModuleOps( |
| 145 | name="upcast_fp8_during_linear_forward", |
| 146 | matcher=lambda model: isinstance(model, LTXModel), |
| 147 | mutator=lambda model: _amend_forward_with_upcast(model, False), |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | class UpcastWithStochasticRounding(ModuleOps): |
| 152 | """ |
| 153 | ModuleOps for upcasting the model's float8_e4m3fn weights and biases to the bfloat16 dtype |
| 154 | and applying stochastic rounding during linear forward. |
| 155 | """ |
| 156 | |
| 157 | def __new__(cls, seed: int = 0): |
| 158 | return super().__new__( |
| 159 | cls, |
| 160 | name="upcast_fp8_during_linear_forward_with_stochastic_rounding", |
| 161 | matcher=lambda model: isinstance(model, LTXModel), |
| 162 | mutator=lambda model: _amend_forward_with_upcast(model, True, seed), |
| 163 | ) |
| 164 |