File size: 3,626 Bytes
f55039b | 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 | """Fused differentiable Qwen operations shared by training and inference."""
from types import MethodType
def fused_rms_norm(module, hidden):
from fla.modules.layernorm import rms_norm
# FLA 0.5.2 accepts bias=None, but its annotation omits None.
return rms_norm(
hidden,
module._jolt_norm_weight,
None, # ty: ignore[invalid-argument-type]
eps=module.eps,
)
def fused_gated_norm(module, hidden, gate):
from fla.modules.fused_norm_gate import rms_norm_gated
# FLA 0.5.2 accepts bias=None, but its annotation omits None.
return rms_norm_gated(
hidden,
gate,
module.weight,
None, # ty: ignore[invalid-argument-type]
eps=module.variance_epsilon,
)
def fused_mlp(module, hidden):
from fla.modules.activations import swiglu
# Keep the projection modules in the graph: they own the trainable LoRA.
return module.down_proj(swiglu(module.gate_proj(hidden), module.up_proj(hidden)))
def enable_fusion(backbone, *, unfused_mlp_layers=()):
from transformers.models.qwen3_5 import modeling_qwen3_5 as qwen
functions = {
qwen.Qwen3_5RMSNorm: fused_rms_norm,
qwen.Qwen3_5RMSNormGated: fused_gated_norm,
qwen.Qwen3_5MLP: fused_mlp,
}
language_layers = backbone.language_model.layers
unfused_mlps = {
id(language_layers[index].mlp) for index in unfused_mlp_layers
}
for module in backbone.modules():
if isinstance(module, qwen.Qwen3_5RMSNorm):
module.register_buffer(
"_jolt_norm_weight", module.weight.detach().float() + 1, persistent=False
)
if type(module) in functions and not (
isinstance(module, qwen.Qwen3_5MLP) and id(module) in unfused_mlps
):
module.forward = MethodType(functions[type(module)], module)
def linear_patch_embedding(module, hidden_states):
"""The patch Conv3d has exactly one output per independent input patch.
Flattening its kernel is the same affine map, without a convolution solver
search for every total patch count. No weights or tokens are added/removed.
"""
from torch.nn import functional
weight = module.proj.weight
patches = hidden_states.reshape(-1, weight[0].numel()).to(weight.dtype)
return functional.linear(patches, weight.flatten(1), module.proj.bias)
def enable_linear_patch_embedding(visual):
module = visual.patch_embed
projection = module.proj
if (
projection.kernel_size != projection.stride
or projection.groups != 1
or any(projection.padding)
or projection.dilation != (1, 1, 1)
):
raise ValueError("Linear patch embedding requires non-overlapping independent patches")
module.forward = MethodType(linear_patch_embedding, module)
def triton_causal_conv1d(hidden_states, weight, bias=None, activation=None, **kwargs):
from fla.modules.conv import causal_conv1d
# FLA's input_guard decorator exposes a Tensor/callable union for this function.
output, _ = causal_conv1d( # ty: ignore[call-non-callable]
hidden_states.transpose(1, 2),
weight=weight,
bias=bias,
activation=activation,
backend="triton",
output_final_state=False,
cu_seqlens=kwargs.get("cu_seq_lens_q"),
)
return output.transpose(1, 2)
def enable_triton_convolution():
"""Explicit process-local override; applies only to full-sequence prefill."""
from transformers.models.qwen3_5 import modeling_qwen3_5
modeling_qwen3_5.causal_conv1d_fn = triton_causal_conv1d
|