From c903a51db04d6ed80d9b5877cab1e1a22c7eafa6 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Tue, 14 Jul 2026 10:53:53 +0000 Subject: [PATCH 1/2] refactor: remove hardcoded DiT LayerNorm paths --- .../networks/ernie_image/infer/post_infer.py | 12 +-- .../ernie_image/weights/post_weights.py | 4 +- .../models/networks/flux2/infer/post_infer.py | 2 +- .../networks/flux2/infer/transformer_infer.py | 10 +-- .../networks/flux2/weights/post_weights.py | 5 +- .../flux2/weights/transformer_weights.py | 11 ++- .../feature_caching/transformer_infer.py | 2 +- .../hunyuan_video/weights/pre_weights.py | 4 +- .../lingbot_video/weights/post_weights.py | 2 +- .../longcat_image/infer/post_infer.py | 2 +- .../longcat_image/infer/transformer_infer.py | 23 ++++-- .../longcat_image/weights/post_weights.py | 5 +- .../weights/transformer_weights.py | 11 ++- .../networks/ltx2/weights/post_weights.py | 4 +- .../qwen_image/weights/post_weights.py | 2 +- .../wan/infer/causvid/transformer_infer.py | 4 +- .../infer/matrix_game2/transformer_infer.py | 2 +- .../networks/wan/infer/s2v/audio_inject.py | 8 +- .../wan/infer/s2v/transformer_infer.py | 10 +-- .../models/networks/wan/infer/s2v/wan_ops.py | 16 ---- .../weights/animate/transformer_weights.py | 4 +- .../wan/weights/audio/transformer_weights.py | 4 +- .../wan/weights/dreamzero/pre_weights.py | 4 +- .../weights/fastwam/transformer_weights.py | 6 +- .../wan/weights/infinitetalk/pre_weights.py | 2 +- .../infinitetalk/transformer_weights.py | 4 +- .../wan/weights/matrix_game2/pre_weights.py | 4 +- .../matrix_game2/transformer_weights.py | 8 +- .../matrix_game3/transformer_weights.py | 6 +- .../networks/wan/weights/motus/_shared.py | 4 +- .../wan/weights/motus/post_weights.py | 2 +- .../networks/wan/weights/pre_weights.py | 4 +- .../wan/weights/s2v/transformer_weights.py | 5 +- .../wan/weights/transformer_weights.py | 8 +- lightx2v/models/networks/worldmirror/model.py | 17 +--- .../weights/transformer_weights.py | 12 ++- .../networks/z_image/weights/post_weights.py | 2 +- .../ops/norm/enflame_gcu/layer_norm.py | 2 +- lightx2v_platform/ops/norm/norm_template.py | 28 +++++-- test_cases/test_layer_norm_weight_routing.py | 82 +++++++++++++++++++ 40 files changed, 219 insertions(+), 128 deletions(-) create mode 100644 test_cases/test_layer_norm_weight_routing.py diff --git a/lightx2v/models/networks/ernie_image/infer/post_infer.py b/lightx2v/models/networks/ernie_image/infer/post_infer.py index 5aace720c..0029c229f 100644 --- a/lightx2v/models/networks/ernie_image/infer/post_infer.py +++ b/lightx2v/models/networks/ernie_image/infer/post_infer.py @@ -1,25 +1,15 @@ -import torch.nn.functional as F - - class ErnieImagePostInfer: def __init__(self, config): self.config = config self.patch_size = config.get("patch_size", 1) self.out_channels = config.get("out_channels", config.get("in_channels", 128)) - self.eps = config.get("eps", 1e-6) def set_scheduler(self, scheduler): self.scheduler = scheduler def infer(self, weights, hidden_states, pre_infer_out): scale, shift = weights.final_norm_linear.apply(pre_infer_out.conditioning).chunk(2, dim=-1) - hidden_states = F.layer_norm( - hidden_states, - (hidden_states.shape[-1],), - weight=None, - bias=None, - eps=self.eps, - ) + hidden_states = weights.final_norm.apply(hidden_states) hidden_states = hidden_states * (1 + scale) + shift patches = weights.final_linear.apply(hidden_states)[: pre_infer_out.image_tokens_len] diff --git a/lightx2v/models/networks/ernie_image/weights/post_weights.py b/lightx2v/models/networks/ernie_image/weights/post_weights.py index ebfcefd93..5021af3dc 100644 --- a/lightx2v/models/networks/ernie_image/weights/post_weights.py +++ b/lightx2v/models/networks/ernie_image/weights/post_weights.py @@ -1,11 +1,13 @@ from lightx2v.common.modules.weight_module import WeightModule -from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER class ErnieImagePostWeights(WeightModule): def __init__(self, config): super().__init__() self.mm_type = config.get("dit_quant_scheme", "Default") + self.layer_norm_type = config.get("layer_norm_type", "torch") + self.add_module("final_norm", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=config.get("eps", 1e-6))) self.add_module( "final_norm_linear", MM_WEIGHT_REGISTER[self.mm_type]( diff --git a/lightx2v/models/networks/flux2/infer/post_infer.py b/lightx2v/models/networks/flux2/infer/post_infer.py index 658335ee4..247a0b78f 100644 --- a/lightx2v/models/networks/flux2/infer/post_infer.py +++ b/lightx2v/models/networks/flux2/infer/post_infer.py @@ -25,7 +25,7 @@ def infer(self, weights, hidden_states, temb): temb_out = weights.norm_out_linear.apply(F.silu(temb)) scale, shift = torch.chunk(temb_out, 2, dim=-1) - hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],)) + hidden_states = weights.norm_out.apply(hidden_states) hidden_states = hidden_states * (1 + scale) + shift output = weights.proj_out.apply(hidden_states) diff --git a/lightx2v/models/networks/flux2/infer/transformer_infer.py b/lightx2v/models/networks/flux2/infer/transformer_infer.py index b4d120990..5236c2e5a 100644 --- a/lightx2v/models/networks/flux2/infer/transformer_infer.py +++ b/lightx2v/models/networks/flux2/infer/transformer_infer.py @@ -63,10 +63,10 @@ def infer_double_stream_block( (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = self._split_double_modulation(temb_mod_img) (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = self._split_double_modulation(temb_mod_txt) - norm_hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],)) + norm_hidden_states = block_weights.norm1.apply(hidden_states) norm_hidden_states = (norm_hidden_states * (1 + scale_msa) + shift_msa).squeeze(0) - norm_encoder_hidden_states = F.layer_norm(encoder_hidden_states, (encoder_hidden_states.shape[-1],)) + norm_encoder_hidden_states = block_weights.norm1_context.apply(encoder_hidden_states) norm_encoder_hidden_states = (norm_encoder_hidden_states * (1 + c_scale_msa) + c_shift_msa).squeeze(0) img_query = block_weights.to_q.apply(norm_hidden_states) @@ -140,7 +140,7 @@ def infer_double_stream_block( img_attn_hook(gated_img_attn) hidden_states = hidden_states + gated_img_attn encoder_hidden_states = encoder_hidden_states + c_gate_msa * txt_attn_output - norm_hidden_states2 = F.layer_norm(hidden_states, (hidden_states.shape[-1],)) + norm_hidden_states2 = block_weights.norm2.apply(hidden_states) norm_hidden_states2 = (norm_hidden_states2 * (1 + scale_mlp) + shift_mlp).squeeze(0) ff_output = block_weights.ff_net_0.apply(norm_hidden_states2) ff_1, ff_2 = ff_output.chunk(2, dim=-1) @@ -148,7 +148,7 @@ def infer_double_stream_block( ff_output = block_weights.ff_net_2.apply(ff_output) hidden_states = hidden_states + gate_mlp * ff_output - norm_encoder_hidden_states2 = F.layer_norm(encoder_hidden_states, (encoder_hidden_states.shape[-1],)) + norm_encoder_hidden_states2 = block_weights.norm2_context.apply(encoder_hidden_states) norm_encoder_hidden_states2 = (norm_encoder_hidden_states2 * (1 + c_scale_mlp) + c_shift_mlp).squeeze(0) context_ff_output = block_weights.ff_context_net_0.apply(norm_encoder_hidden_states2) ctx_ff_1, ctx_ff_2 = context_ff_output.chunk(2, dim=-1) @@ -179,7 +179,7 @@ def infer_single_stream_block( shift_msa, scale_msa, gate_msa = self._split_single_modulation(temb_mod) - norm_combined = F.layer_norm(hidden_states, (hidden_states.shape[-1],)) + norm_combined = block_weights.norm.apply(hidden_states) norm_combined = (norm_combined * (1 + scale_msa) + shift_msa).squeeze(0) hidden_states_proj = block_weights.to_qkv_mlp_proj.apply(norm_combined) diff --git a/lightx2v/models/networks/flux2/weights/post_weights.py b/lightx2v/models/networks/flux2/weights/post_weights.py index 6a3bf0c5d..0c90f21d4 100644 --- a/lightx2v/models/networks/flux2/weights/post_weights.py +++ b/lightx2v/models/networks/flux2/weights/post_weights.py @@ -1,5 +1,5 @@ from lightx2v.common.modules.weight_module import WeightModule -from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER class Flux2PostWeights(WeightModule): @@ -12,6 +12,9 @@ def __init__(self, config): self.out_channels = config.get("transformer_in_channels", config.get("in_channels", 64)) self.patch_size = config.get("patch_size", 1) self.mm_type = config.get("dit_quant_scheme", "Default") + self.layer_norm_type = config.get("layer_norm_type", "torch") + + self.add_module("norm_out", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) self.add_module( "norm_out_linear", diff --git a/lightx2v/models/networks/flux2/weights/transformer_weights.py b/lightx2v/models/networks/flux2/weights/transformer_weights.py index da3ed9141..2777c357f 100644 --- a/lightx2v/models/networks/flux2/weights/transformer_weights.py +++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py @@ -1,5 +1,5 @@ from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER class Flux2DoubleBlockWeights(WeightModule): @@ -11,11 +11,17 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe self.block_idx = block_idx self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"] self.mm_type = config.get("dit_quant_scheme", "Default") + self.layer_norm_type = config.get("layer_norm_type", "torch") self.rms_norm_type = config.get("rms_norm_type", "torch") self.attn_type = config.get("attn_type", "flash_attn3") p = f"transformer_blocks.{self.block_idx}" + self.add_module("norm1", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + self.add_module("norm1_context", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + self.add_module("norm2", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + self.add_module("norm2_context", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + self.add_module( "to_q", MM_WEIGHT_REGISTER[self.mm_type]( @@ -189,11 +195,14 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe self.block_idx = block_idx self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"] self.mm_type = config.get("dit_quant_scheme", "Default") + self.layer_norm_type = config.get("layer_norm_type", "torch") self.rms_norm_type = config.get("rms_norm_type", "torch") self.attn_type = config.get("attn_type", "flash_attn3") p = f"single_transformer_blocks.{self.block_idx}" + self.add_module("norm", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + self.add_module( "to_qkv_mlp_proj", MM_WEIGHT_REGISTER[self.mm_type]( diff --git a/lightx2v/models/networks/hunyuan_video/infer/feature_caching/transformer_infer.py b/lightx2v/models/networks/hunyuan_video/infer/feature_caching/transformer_infer.py index 56041ec8b..d35362f57 100755 --- a/lightx2v/models/networks/hunyuan_video/infer/feature_caching/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_video/infer/feature_caching/transformer_infer.py @@ -151,7 +151,7 @@ def calculate_should_calc(self, img, vec, block): img_mod1_shift, img_mod1_scale, _, _, _, _ = img_mod_layer.apply(vec_).chunk(6, dim=-1) inp = inp.squeeze(0) - normed_inp = torch.nn.functional.layer_norm(inp, (inp.shape[1],), None, None, 1e-6) + normed_inp = block.img_branch.img_norm1.apply(inp) modulated_inp = normed_inp * (1 + img_mod1_scale) + img_mod1_shift del normed_inp, inp, vec_ diff --git a/lightx2v/models/networks/hunyuan_video/weights/pre_weights.py b/lightx2v/models/networks/hunyuan_video/weights/pre_weights.py index e80298b4b..9f9c8f943 100644 --- a/lightx2v/models/networks/hunyuan_video/weights/pre_weights.py +++ b/lightx2v/models/networks/hunyuan_video/weights/pre_weights.py @@ -120,7 +120,7 @@ def __init__(self, block_idx, mm_type, config, block_prefix): self.mm_type = mm_type self.add_module( "norm1", - LN_WEIGHT_REGISTER["torch"](f"{block_prefix}.{block_idx}.norm1.weight", f"{block_prefix}.{block_idx}.norm1.bias"), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](f"{block_prefix}.{block_idx}.norm1.weight", f"{block_prefix}.{block_idx}.norm1.bias"), ) self.add_module( "self_attn_qkv", @@ -134,7 +134,7 @@ def __init__(self, block_idx, mm_type, config, block_prefix): self.add_module("self_attention", attention_weights_cls()) self.add_module( "norm2", - LN_WEIGHT_REGISTER["torch"](f"{block_prefix}.{block_idx}.norm2.weight", f"{block_prefix}.{block_idx}.norm2.bias"), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](f"{block_prefix}.{block_idx}.norm2.weight", f"{block_prefix}.{block_idx}.norm2.bias"), ) self.add_module( "mlp_fc1", diff --git a/lightx2v/models/networks/lingbot_video/weights/post_weights.py b/lightx2v/models/networks/lingbot_video/weights/post_weights.py index 2c22aae49..9fe30f561 100644 --- a/lightx2v/models/networks/lingbot_video/weights/post_weights.py +++ b/lightx2v/models/networks/lingbot_video/weights/post_weights.py @@ -5,6 +5,6 @@ class LingBotVideoPostWeights(WeightModule): def __init__(self, config): super().__init__() - self.add_module("norm_out", LN_WEIGHT_REGISTER["torch"](eps=config.get("norm_eps", 1e-6))) + self.add_module("norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](eps=config.get("norm_eps", 1e-6))) self.add_module("norm_out_modulation", MM_WEIGHT_REGISTER["Default-ForceFp32"]("norm_out_modulation.1.weight", "norm_out_modulation.1.bias")) self.add_module("proj_out", MM_WEIGHT_REGISTER["Default"]("proj_out.weight", "proj_out.bias")) diff --git a/lightx2v/models/networks/longcat_image/infer/post_infer.py b/lightx2v/models/networks/longcat_image/infer/post_infer.py index 372640ce3..78b0c3fea 100755 --- a/lightx2v/models/networks/longcat_image/infer/post_infer.py +++ b/lightx2v/models/networks/longcat_image/infer/post_infer.py @@ -30,7 +30,7 @@ def infer(self, weights, hidden_states, temb): scale, shift = torch.chunk(temb_out, 2, dim=-1) # Layer norm (no learnable params) + modulation - hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],)) + hidden_states = weights.norm_out.apply(hidden_states) hidden_states = hidden_states * (1 + scale) + shift # Final projection diff --git a/lightx2v/models/networks/longcat_image/infer/transformer_infer.py b/lightx2v/models/networks/longcat_image/infer/transformer_infer.py index b8e70d36d..e541754bc 100755 --- a/lightx2v/models/networks/longcat_image/infer/transformer_infer.py +++ b/lightx2v/models/networks/longcat_image/infer/transformer_infer.py @@ -40,7 +40,7 @@ def __init__(self, config): def set_scheduler(self, scheduler): self.scheduler = scheduler - def _ada_layer_norm_zero(self, hidden_states, temb, norm_linear): + def _ada_layer_norm_zero(self, hidden_states, temb, norm, norm_linear): """AdaLayerNormZero: compute shift, scale, gate from temb. For double-stream blocks: returns (norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp) @@ -51,12 +51,12 @@ def _ada_layer_norm_zero(self, hidden_states, temb, norm_linear): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) # Apply layer norm and modulation - norm_hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],)) + norm_hidden_states = norm.apply(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa return norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp - def _ada_layer_norm_zero_single(self, hidden_states, temb, norm_linear): + def _ada_layer_norm_zero_single(self, hidden_states, temb, norm, norm_linear): """AdaLayerNormZeroSingle: for single-stream blocks. Returns (norm_hidden_states, gate) @@ -67,7 +67,7 @@ def _ada_layer_norm_zero_single(self, hidden_states, temb, norm_linear): shift, scale, gate = emb.chunk(3, dim=-1) # Apply layer norm and modulation - norm_hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],)) + norm_hidden_states = norm.apply(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale) + shift return norm_hidden_states, gate @@ -100,10 +100,15 @@ def infer_double_stream_block( head_dim = self.config["attention_head_dim"] # ===== Image stream: norm1 ===== - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._ada_layer_norm_zero(hidden_states, temb, block_weights.norm1_linear) + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._ada_layer_norm_zero(hidden_states, temb, block_weights.norm1, block_weights.norm1_linear) # ===== Text stream: norm1_context ===== - norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self._ada_layer_norm_zero(encoder_hidden_states, temb, block_weights.norm1_context_linear) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self._ada_layer_norm_zero( + encoder_hidden_states, + temb, + block_weights.norm1_context, + block_weights.norm1_context_linear, + ) # ===== Attention projections ===== # Image stream QKV @@ -186,7 +191,7 @@ def infer_double_stream_block( # ===== FFN for image stream ===== # Layer norm without learnable parameters (LongCat/Flux architecture) - norm_hidden_states2 = F.layer_norm(hidden_states, (hidden_states.shape[-1],)) + norm_hidden_states2 = block_weights.norm2.apply(hidden_states) norm_hidden_states2 = norm_hidden_states2 * (1 + scale_mlp) + shift_mlp ff_output = block_weights.ff_net_0_proj.apply(norm_hidden_states2) ff_output = F.gelu(ff_output, approximate="tanh") @@ -195,7 +200,7 @@ def infer_double_stream_block( # ===== FFN for text stream ===== # Layer norm without learnable parameters (LongCat/Flux architecture) - norm_encoder_hidden_states2 = F.layer_norm(encoder_hidden_states, (encoder_hidden_states.shape[-1],)) + norm_encoder_hidden_states2 = block_weights.norm2_context.apply(encoder_hidden_states) norm_encoder_hidden_states2 = norm_encoder_hidden_states2 * (1 + c_scale_mlp) + c_shift_mlp context_ff_output = block_weights.ff_context_net_0_proj.apply(norm_encoder_hidden_states2) context_ff_output = F.gelu(context_ff_output, approximate="tanh") @@ -242,7 +247,7 @@ def infer_single_stream_block( residual = combined # AdaLayerNormZeroSingle - norm_combined, gate = self._ada_layer_norm_zero_single(combined, temb, block_weights.norm_linear) + norm_combined, gate = self._ada_layer_norm_zero_single(combined, temb, block_weights.norm, block_weights.norm_linear) # MLP branch mlp_hidden_states = block_weights.proj_mlp.apply(norm_combined) diff --git a/lightx2v/models/networks/longcat_image/weights/post_weights.py b/lightx2v/models/networks/longcat_image/weights/post_weights.py index a690edef9..a6654abe3 100755 --- a/lightx2v/models/networks/longcat_image/weights/post_weights.py +++ b/lightx2v/models/networks/longcat_image/weights/post_weights.py @@ -1,5 +1,5 @@ from lightx2v.common.modules.weight_module import WeightModule -from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER class LongCatImagePostWeights(WeightModule): @@ -13,6 +13,9 @@ def __init__(self, config): self.out_channels = config.get("transformer_in_channels", config.get("in_channels", 64)) self.patch_size = config.get("patch_size", 1) self.mm_type = config.get("dit_quant_scheme", "Default") + self.layer_norm_type = config.get("layer_norm_type", "torch") + + self.add_module("norm_out", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) # norm_out (AdaLayerNormContinuous) self.add_module( diff --git a/lightx2v/models/networks/longcat_image/weights/transformer_weights.py b/lightx2v/models/networks/longcat_image/weights/transformer_weights.py index 7c081d964..6307ab50e 100755 --- a/lightx2v/models/networks/longcat_image/weights/transformer_weights.py +++ b/lightx2v/models/networks/longcat_image/weights/transformer_weights.py @@ -1,5 +1,5 @@ from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER class LongCatImageDoubleBlockWeights(WeightModule): @@ -11,11 +11,17 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe self.block_idx = block_idx self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"] self.mm_type = config.get("dit_quant_scheme", "Default") + self.layer_norm_type = config.get("layer_norm_type", "torch") self.rms_norm_type = config.get("rms_norm_type", "torch") self.attn_type = config.get("attn_type", "flash_attn3") p = f"transformer_blocks.{self.block_idx}" + self.add_module("norm1", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + self.add_module("norm1_context", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + self.add_module("norm2", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + self.add_module("norm2_context", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + # Image stream norm1 (AdaLayerNormZero) self.add_module( "norm1_linear", @@ -216,11 +222,14 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe self.block_idx = block_idx self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"] self.mm_type = config.get("dit_quant_scheme", "Default") + self.layer_norm_type = config.get("layer_norm_type", "torch") self.rms_norm_type = config.get("rms_norm_type", "torch") self.attn_type = config.get("attn_type", "flash_attn3") p = f"single_transformer_blocks.{self.block_idx}" + self.add_module("norm", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5)) + # AdaLayerNormZeroSingle self.add_module( "norm_linear", diff --git a/lightx2v/models/networks/ltx2/weights/post_weights.py b/lightx2v/models/networks/ltx2/weights/post_weights.py index bba0cfd83..27db1e439 100755 --- a/lightx2v/models/networks/ltx2/weights/post_weights.py +++ b/lightx2v/models/networks/ltx2/weights/post_weights.py @@ -45,7 +45,7 @@ def __init__(self, config): "model.diffusion_model.scale_shift_table", ), ) - self.add_module("norm_out", LN_WEIGHT_REGISTER["torch"]()) + self.add_module("norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) self.add_module( "proj_out", MM_WEIGHT_REGISTER["Default"]( @@ -60,7 +60,7 @@ def __init__(self, config): "model.diffusion_model.audio_scale_shift_table", ), ) - self.add_module("audio_norm_out", LN_WEIGHT_REGISTER["torch"]()) + self.add_module("audio_norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) self.add_module( "audio_proj_out", MM_WEIGHT_REGISTER["Default"]( diff --git a/lightx2v/models/networks/qwen_image/weights/post_weights.py b/lightx2v/models/networks/qwen_image/weights/post_weights.py index 4fce38e87..fe545e30f 100755 --- a/lightx2v/models/networks/qwen_image/weights/post_weights.py +++ b/lightx2v/models/networks/qwen_image/weights/post_weights.py @@ -18,7 +18,7 @@ def __init__(self, config): "norm_out.linear.bias", ), ) - self.add_module("norm_out", LN_WEIGHT_REGISTER["torch"](eps=1e-6)) + self.add_module("norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](eps=1e-6)) # proj_out self.add_module( diff --git a/lightx2v/models/networks/wan/infer/causvid/transformer_infer.py b/lightx2v/models/networks/wan/infer/causvid/transformer_infer.py index 236f8075e..aa62a23c7 100755 --- a/lightx2v/models/networks/wan/infer/causvid/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/causvid/transformer_infer.py @@ -96,7 +96,7 @@ def _infer_without_offload(self, weights, grid_sizes, embed, x, embed0, seq_lens return x def infer_self_attn(self, weights, grid_sizes, embed, x, embed0, seq_lens, freqs, context, block_idx, kv_start, kv_end): - norm1_out = torch.nn.functional.layer_norm(x, (x.shape[1],), None, None, 1e-6) + norm1_out = weights.norm1.apply(x) norm1_out = (norm1_out * (1 + embed0[1]) + embed0[0]).squeeze(0) s, n, d = *norm1_out.shape[:1], self.num_heads, self.head_dim @@ -197,7 +197,7 @@ def infer_cross_attn(self, weights, x, context, block_idx): return x def infer_ffn(self, weights, x, embed0): - norm2_out = torch.nn.functional.layer_norm(x, (x.shape[1],), None, None, 1e-6) + norm2_out = weights.norm2.apply(x) y = weights.ffn_0.apply(norm2_out * (1 + embed0[4].squeeze(0)) + embed0[3].squeeze(0)) y = torch.nn.functional.gelu(y, approximate="tanh") y = weights.ffn_2.apply(y) diff --git a/lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py b/lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py index 48e4ca27a..950b78812 100755 --- a/lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py @@ -168,7 +168,7 @@ def _infer_action_model_with_kvcache( group_mouse = torch.nn.functional.linear(group_mouse, phase.mouse_mlp_0.weight.T, phase.mouse_mlp_0.bias) group_mouse = torch.nn.functional.gelu(group_mouse, approximate="tanh") group_mouse = torch.nn.functional.linear(group_mouse, phase.mouse_mlp_2.weight.T, phase.mouse_mlp_2.bias) - group_mouse = torch.nn.functional.layer_norm(group_mouse, (group_mouse.shape[-1],), phase.mouse_mlp_3.weight.T, phase.mouse_mlp_3.bias, 1e-5) + group_mouse = phase.mouse_mlp_3.apply(group_mouse) mouse_qkv = torch.nn.functional.linear(group_mouse, phase.t_qkv.weight.T) diff --git a/lightx2v/models/networks/wan/infer/s2v/audio_inject.py b/lightx2v/models/networks/wan/infer/s2v/audio_inject.py index 1026e8491..ec72eb281 100644 --- a/lightx2v/models/networks/wan/infer/s2v/audio_inject.py +++ b/lightx2v/models/networks/wan/infer/s2v/audio_inject.py @@ -6,12 +6,12 @@ from lightx2v_platform.base.global_var import AI_DEVICE -def _apply_adain(x, temb, adain_linear): +def _apply_adain(x, temb, adain_linear, adain_norm): """Mirror diffusers AdaLayerNorm(chunk_dim=1) used by Wan audio_injector.""" temb = F.silu(temb) emb = mm_weight_autocast_nd(adain_linear, temb) shift, scale = emb.chunk(2, dim=-1) - x = F.layer_norm(x, (x.shape[-1],), weight=None, bias=None, eps=1e-5) + x = adain_norm.apply(x) return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) @@ -32,9 +32,9 @@ def apply_audio_inject(block, hidden_states, pre_infer_out, config): with torch.amp.autocast(str(AI_DEVICE), dtype=torch.bfloat16): if config.get("enable_adain", False) and config.get("adain_mode") == "attn_norm": audio_emb_global = rearrange(pre_infer_out.audio_emb_global, "b t n c -> (b t) n c", t=num_frames) - attn_hidden_states = _apply_adain(input_hidden_states, audio_emb_global[:, 0], inj.adain_linear) + attn_hidden_states = _apply_adain(input_hidden_states, audio_emb_global[:, 0], inj.adain_linear, inj.adain_norm) else: - attn_hidden_states = F.layer_norm(input_hidden_states, (input_hidden_states.shape[-1],), weight=None, bias=None, eps=1e-6) + attn_hidden_states = inj.attn_norm.apply(input_hidden_states) residual_out = cross_attn_forward(inj, attn_hidden_states, audio_emb, None, n, d) residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames) hidden_states[:, :original_seq_len] = hidden_states[:, :original_seq_len] + residual_out diff --git a/lightx2v/models/networks/wan/infer/s2v/transformer_infer.py b/lightx2v/models/networks/wan/infer/s2v/transformer_infer.py index 4f16fc443..724c457d5 100644 --- a/lightx2v/models/networks/wan/infer/s2v/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/s2v/transformer_infer.py @@ -10,8 +10,6 @@ s2v_self_attn_forward, segment_gate_bld, segment_modulate_bld, - wan_layer_norm, - wan_layer_norm_float, wan_rms_norm, ) from lightx2v.models.networks.wan.infer.transformer_infer import WanTransformerInfer @@ -42,7 +40,7 @@ def infer_non_blocks(self, weights, x, e, original_seq_len=None): modulation = weights.head_modulation.tensor with torch.amp.autocast(AI_DEVICE, dtype=torch.float32): e = (modulation + e.unsqueeze(1)).chunk(2, dim=1) - norm_x = wan_layer_norm(weights.norm, x).float() + norm_x = weights.norm.apply(x).float() x = norm_x * (1 + e[1]) + e[0] x = mm_weight_fp32_nd(weights.head, x) if x.dim() == 3: @@ -133,7 +131,7 @@ def infer_s2v_block(self, block, x, pre_infer_out): seq_lens = torch.tensor([x.size(1)], dtype=torch.long, device=x.device) # norm1 + modulate (Wan 212-217) - norm_x = wan_layer_norm_float(phase0.norm1, x) + norm_x = phase0.norm1.apply(x).float() norm_x = segment_modulate_bld(norm_x, e[0], e[1], seg_idx) # self-attention (Wan 219-225) @@ -145,9 +143,9 @@ def infer_s2v_block(self, block, x, pre_infer_out): # cross-attention & ffn (Wan cross_attn_ffn 228-243) context = pre_infer_out.context ctx_lens = pre_infer_out.s2v_extra.get("context_lens") - x = x + cross_attn_forward(phase1, wan_layer_norm(phase1.norm3, x), context, ctx_lens, self.num_heads, self.head_dim) + x = x + cross_attn_forward(phase1, phase1.norm3.apply(x), context, ctx_lens, self.num_heads, self.head_dim) - norm2_x = wan_layer_norm_float(phase2.norm2, x) + norm2_x = phase2.norm2.apply(x).float() norm2_x = segment_modulate_bld(norm2_x, e[3], e[4], seg_idx) y = mm_weight_autocast_nd(phase2.ffn_0, norm2_x) y = torch.nn.functional.gelu(y, approximate="tanh") diff --git a/lightx2v/models/networks/wan/infer/s2v/wan_ops.py b/lightx2v/models/networks/wan/infer/s2v/wan_ops.py index 0602abdce..607762178 100644 --- a/lightx2v/models/networks/wan/infer/s2v/wan_ops.py +++ b/lightx2v/models/networks/wan/infer/s2v/wan_ops.py @@ -38,22 +38,6 @@ def mm_weight_autocast_nd(linear, x, autocast_dtype=torch.bfloat16): return out.view(*shape[:-1], -1) -def wan_layer_norm(ln_weight, x, force_float=False): - weight = ln_weight._get_actual_weight() if ln_weight.weight is not None else None - bias = ln_weight._get_actual_bias() if getattr(ln_weight, "bias", None) is not None else None - if weight is not None: - weight = weight.float() - if bias is not None: - bias = bias.float() - out = torch.nn.functional.layer_norm(x.float(), (x.shape[-1],), weight, bias, ln_weight.eps) - out = out.to(x.dtype) - return out.float() if force_float else out - - -def wan_layer_norm_float(ln_weight, x): - return wan_layer_norm(ln_weight, x, force_float=True) - - def wan_rms_norm(rms_weight, x): w = rms_weight._get_actual_weight() xf = x.float() diff --git a/lightx2v/models/networks/wan/weights/animate/transformer_weights.py b/lightx2v/models/networks/wan/weights/animate/transformer_weights.py index 88869a15d..449c02bf8 100755 --- a/lightx2v/models/networks/wan/weights/animate/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/animate/transformer_weights.py @@ -116,12 +116,12 @@ def __init__(self, config, block_index, block_prefix, mm_type, create_cuda_buffe self.add_module( "pre_norm_feat", - LN_WEIGHT_REGISTER["torch"](), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](), ) self.add_module( "pre_norm_motion", - LN_WEIGHT_REGISTER["torch"](), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](), ) self.add_module("adapter_attn", ATTN_WEIGHT_REGISTER[config["adapter_attn_type"]]()) diff --git a/lightx2v/models/networks/wan/weights/audio/transformer_weights.py b/lightx2v/models/networks/wan/weights/audio/transformer_weights.py index 528067880..ed35db645 100755 --- a/lightx2v/models/networks/wan/weights/audio/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/audio/transformer_weights.py @@ -134,7 +134,7 @@ def __init__(self, block_index, block_prefix, task, mm_type, config, create_cuda self.add_module( "norm_kv", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( f"{block_prefix}.{block_index}.norm_kv.weight", f"{block_prefix}.{block_index}.norm_kv.bias", create_cuda_buffer, @@ -146,7 +146,7 @@ def __init__(self, block_index, block_prefix, task, mm_type, config, create_cuda self.add_module( "norm_q", - LN_WEIGHT_REGISTER["torch"](), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](), ) self.add_module( diff --git a/lightx2v/models/networks/wan/weights/dreamzero/pre_weights.py b/lightx2v/models/networks/wan/weights/dreamzero/pre_weights.py index e87ce3499..cb4528c81 100644 --- a/lightx2v/models/networks/wan/weights/dreamzero/pre_weights.py +++ b/lightx2v/models/networks/wan/weights/dreamzero/pre_weights.py @@ -84,7 +84,7 @@ def __init__(self, config): ) self.add_module( "proj_0", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( "img_emb.proj.0.weight", "img_emb.proj.0.bias", eps=1e-5, @@ -109,7 +109,7 @@ def __init__(self, config): ) self.add_module( "proj_4", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( "img_emb.proj.4.weight", "img_emb.proj.4.bias", eps=1e-5, diff --git a/lightx2v/models/networks/wan/weights/fastwam/transformer_weights.py b/lightx2v/models/networks/wan/weights/fastwam/transformer_weights.py index d7385b3a9..bd9546b27 100644 --- a/lightx2v/models/networks/wan/weights/fastwam/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/fastwam/transformer_weights.py @@ -10,7 +10,7 @@ def __init__(self, prefix, block_index, config): eps = float(config.get("eps", 1e-6)) self.add_module("modulation", TENSOR_REGISTER["Default"](f"{block}.modulation")) - self.add_module("norm1", LN_WEIGHT_REGISTER["torch"](eps=eps)) + self.add_module("norm1", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](eps=eps)) self.add_module("q", MM_WEIGHT_REGISTER["Default"](f"{block}.self_attn.q.weight", f"{block}.self_attn.q.bias")) self.add_module("k", MM_WEIGHT_REGISTER["Default"](f"{block}.self_attn.k.weight", f"{block}.self_attn.k.bias")) self.add_module("v", MM_WEIGHT_REGISTER["Default"](f"{block}.self_attn.v.weight", f"{block}.self_attn.v.bias")) @@ -27,7 +27,7 @@ def __init__(self, prefix, block_index, config): rms_type = config.get("rms_norm_type", "torch") eps = float(config.get("eps", 1e-6)) - self.add_module("norm3", LN_WEIGHT_REGISTER["torch"](f"{block}.norm3.weight", f"{block}.norm3.bias", eps=eps)) + self.add_module("norm3", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](f"{block}.norm3.weight", f"{block}.norm3.bias", eps=eps)) self.add_module("q", MM_WEIGHT_REGISTER["Default"](f"{block}.cross_attn.q.weight", f"{block}.cross_attn.q.bias")) self.add_module("k", MM_WEIGHT_REGISTER["Default"](f"{block}.cross_attn.k.weight", f"{block}.cross_attn.k.bias")) self.add_module("v", MM_WEIGHT_REGISTER["Default"](f"{block}.cross_attn.v.weight", f"{block}.cross_attn.v.bias")) @@ -43,7 +43,7 @@ def __init__(self, prefix, block_index, config): block = f"{prefix}.blocks.{block_index}" eps = float(config.get("eps", 1e-6)) - self.add_module("norm2", LN_WEIGHT_REGISTER["torch"](eps=eps)) + self.add_module("norm2", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](eps=eps)) self.add_module("fc0", MM_WEIGHT_REGISTER["Default"](f"{block}.ffn.0.weight", f"{block}.ffn.0.bias")) self.add_module("fc2", MM_WEIGHT_REGISTER["Default"](f"{block}.ffn.2.weight", f"{block}.ffn.2.bias")) diff --git a/lightx2v/models/networks/wan/weights/infinitetalk/pre_weights.py b/lightx2v/models/networks/wan/weights/infinitetalk/pre_weights.py index 8e0119a3c..0c160dfca 100644 --- a/lightx2v/models/networks/wan/weights/infinitetalk/pre_weights.py +++ b/lightx2v/models/networks/wan/weights/infinitetalk/pre_weights.py @@ -24,5 +24,5 @@ def __init__(self, config): if config.get("norm_output_audio", True): self.add_module( "audio_proj_norm", - LN_WEIGHT_REGISTER["torch"]("audio_proj.norm.weight", "audio_proj.norm.bias"), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]("audio_proj.norm.weight", "audio_proj.norm.bias"), ) diff --git a/lightx2v/models/networks/wan/weights/infinitetalk/transformer_weights.py b/lightx2v/models/networks/wan/weights/infinitetalk/transformer_weights.py index 303b68ca5..06e127b06 100755 --- a/lightx2v/models/networks/wan/weights/infinitetalk/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/infinitetalk/transformer_weights.py @@ -29,7 +29,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): self.register_offload_buffers(config, lazy_load_path, lora_path) self.add_module("blocks", self.blocks) - self.register_parameter("norm", LN_WEIGHT_REGISTER["torch"]()) + self.register_parameter("norm", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) self.add_module( "head", MM_WEIGHT_REGISTER["Default"]( @@ -186,7 +186,7 @@ def __init__( self.config = config self.add_module( "norm_x", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( f"{block_prefix}.{block_index}.norm_x.weight", f"{block_prefix}.{block_index}.norm_x.bias", create_cuda_buffer, diff --git a/lightx2v/models/networks/wan/weights/matrix_game2/pre_weights.py b/lightx2v/models/networks/wan/weights/matrix_game2/pre_weights.py index 574194486..39c02c256 100644 --- a/lightx2v/models/networks/wan/weights/matrix_game2/pre_weights.py +++ b/lightx2v/models/networks/wan/weights/matrix_game2/pre_weights.py @@ -34,7 +34,7 @@ def __init__(self, config): # img_emb self.add_module( "img_emb_0", - LN_WEIGHT_REGISTER["torch"]("img_emb.proj.0.weight", "img_emb.proj.0.bias", eps=1e-5), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]("img_emb.proj.0.weight", "img_emb.proj.0.bias", eps=1e-5), ) self.add_module( "img_emb_1", @@ -46,5 +46,5 @@ def __init__(self, config): ) self.add_module( "img_emb_4", - LN_WEIGHT_REGISTER["torch"]("img_emb.proj.4.weight", "img_emb.proj.4.bias", eps=1e-5), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]("img_emb.proj.4.weight", "img_emb.proj.4.bias", eps=1e-5), ) diff --git a/lightx2v/models/networks/wan/weights/matrix_game2/transformer_weights.py b/lightx2v/models/networks/wan/weights/matrix_game2/transformer_weights.py index fc006fca7..230657140 100755 --- a/lightx2v/models/networks/wan/weights/matrix_game2/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/matrix_game2/transformer_weights.py @@ -34,7 +34,7 @@ def __init__(self, config): self.add_module("blocks", self.blocks) # non blocks weights - self.register_parameter("norm", LN_WEIGHT_REGISTER["torch"]()) + self.register_parameter("norm", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) self.add_module("head", MM_WEIGHT_REGISTER["Default"]("head.head.weight", "head.head.bias")) self.register_parameter("head_modulation", TENSOR_REGISTER["Default"]("head.modulation")) @@ -177,10 +177,10 @@ def __init__(self, block_index, block_prefix, task, mm_type, config): ) self.add_module( "mouse_mlp_3", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( f"{block_prefix}.{self.block_index}.action_model.mouse_mlp.3.weight", f"{block_prefix}.{self.block_index}.action_model.mouse_mlp.3.bias", - eps=1e-6, + eps=1e-5, ), ) @@ -200,7 +200,7 @@ def __init__(self, block_index, block_prefix, task, mm_type, config): self.add_module( "norm3", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( f"{block_prefix}.{self.block_index}.norm3.weight", f"{block_prefix}.{self.block_index}.norm3.bias", ), diff --git a/lightx2v/models/networks/wan/weights/matrix_game3/transformer_weights.py b/lightx2v/models/networks/wan/weights/matrix_game3/transformer_weights.py index 40e91b8f3..e42001ace 100644 --- a/lightx2v/models/networks/wan/weights/matrix_game3/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/matrix_game3/transformer_weights.py @@ -43,7 +43,7 @@ def __init__(self, config): self.add_module("blocks", self.blocks) # Non-block weights (head) - self.register_parameter("norm", LN_WEIGHT_REGISTER["torch"]()) + self.register_parameter("norm", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) self.add_module("head", MM_WEIGHT_REGISTER["Default"]("head.head.weight", "head.head.bias")) self.register_parameter("head_modulation", TENSOR_REGISTER["Default"]("head.modulation")) @@ -154,7 +154,7 @@ def __init__(self, block_index, block_prefix, task, mm_type, config): # norm3 with elementwise_affine=True for cross_attn_norm self.add_module( "norm3", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( f"{block_prefix}.{block_index}.norm3.weight", f"{block_prefix}.{block_index}.norm3.bias", ), @@ -255,7 +255,7 @@ def __init__(self, block_index, block_prefix, task, mm_type, config): ) self.add_module( "mouse_mlp_3", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( f"{block_prefix}.{block_index}.action_model.mouse_mlp.3.weight", f"{block_prefix}.{block_index}.action_model.mouse_mlp.3.bias", eps=1e-5, diff --git a/lightx2v/models/networks/wan/weights/motus/_shared.py b/lightx2v/models/networks/wan/weights/motus/_shared.py index 0efadef13..ae9f8eef8 100644 --- a/lightx2v/models/networks/wan/weights/motus/_shared.py +++ b/lightx2v/models/networks/wan/weights/motus/_shared.py @@ -249,8 +249,8 @@ def __init__( super().__init__() self.block_idx = block_idx - self.add_module("norm1", LN_WEIGHT_REGISTER["torch"]()) - self.add_module("norm2", LN_WEIGHT_REGISTER["torch"]()) + self.add_module("norm1", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) + self.add_module("norm2", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) self.add_module( f"{attr_prefix}_qkv", PackedQKVWeights( diff --git a/lightx2v/models/networks/wan/weights/motus/post_weights.py b/lightx2v/models/networks/wan/weights/motus/post_weights.py index 8b8bf5168..8adb5d165 100644 --- a/lightx2v/models/networks/wan/weights/motus/post_weights.py +++ b/lightx2v/models/networks/wan/weights/motus/post_weights.py @@ -9,7 +9,7 @@ class MotusActionPostWeights(WeightModule): def __init__(self, config, action_config): super().__init__() self.action_config = action_config - self.add_module("norm", LN_WEIGHT_REGISTER["torch"]()) + self.add_module("norm", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) self.add_module( "action_head", SequentialLinearWeights( diff --git a/lightx2v/models/networks/wan/weights/pre_weights.py b/lightx2v/models/networks/wan/weights/pre_weights.py index 5d450d3fb..c71d2aa88 100755 --- a/lightx2v/models/networks/wan/weights/pre_weights.py +++ b/lightx2v/models/networks/wan/weights/pre_weights.py @@ -100,7 +100,7 @@ def __init__(self, config): if config["task"] in ["i2v", "flf2v", "animate", "s2v", "rs2v"] and config.get("use_image_encoder", True): self.add_module( "proj_0", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( "img_emb.proj.0.weight", "img_emb.proj.0.bias", lora_prefix="diffusion_model.img_emb", @@ -124,7 +124,7 @@ def __init__(self, config): ) self.add_module( "proj_4", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( "img_emb.proj.4.weight", "img_emb.proj.4.bias", lora_prefix="diffusion_model.img_emb", diff --git a/lightx2v/models/networks/wan/weights/s2v/transformer_weights.py b/lightx2v/models/networks/wan/weights/s2v/transformer_weights.py index b8058c70d..3b74cb8f4 100644 --- a/lightx2v/models/networks/wan/weights/s2v/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/s2v/transformer_weights.py @@ -1,12 +1,15 @@ from lightx2v.common.modules.weight_module import WeightModule from lightx2v.models.networks.wan.weights.transformer_weights import WanTransformerWeights -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER class WanS2VAudioInjectPhase(WeightModule): def __init__(self, injector_index, mm_type, config): super().__init__() prefix = f"audio_injector.injector.{injector_index}" + layer_norm_type = config.get("layer_norm_type", "torch") + self.add_module("adain_norm", LN_WEIGHT_REGISTER[layer_norm_type](eps=1e-5)) + self.add_module("attn_norm", LN_WEIGHT_REGISTER[layer_norm_type](eps=1e-6)) self.add_module( "norm_q", RMS_WEIGHT_REGISTER["sgl-kernel"](f"{prefix}.norm_q.weight"), diff --git a/lightx2v/models/networks/wan/weights/transformer_weights.py b/lightx2v/models/networks/wan/weights/transformer_weights.py index 23cb4753c..b22454894 100755 --- a/lightx2v/models/networks/wan/weights/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/transformer_weights.py @@ -41,7 +41,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): self.add_module("blocks", self.blocks) # non blocks weights - self.register_parameter("norm", LN_WEIGHT_REGISTER["torch"]()) + self.register_parameter("norm", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]()) self.add_module( "head", MM_WEIGHT_REGISTER["Default"]( @@ -253,7 +253,7 @@ def __init__( self.add_module( "norm1", - LN_WEIGHT_REGISTER["torch"](), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](), ) self.add_module( @@ -480,7 +480,7 @@ def __init__( self.add_module( "norm3", - LN_WEIGHT_REGISTER["torch"]( + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( f"{block_prefix}.{self.block_index}.norm3.weight", f"{block_prefix}.{self.block_index}.norm3.bias", create_cuda_buffer, @@ -636,7 +636,7 @@ def __init__( self.add_module( "norm2", - LN_WEIGHT_REGISTER["torch"](), + LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](), ) self.add_module( diff --git a/lightx2v/models/networks/worldmirror/model.py b/lightx2v/models/networks/worldmirror/model.py index e3f61def5..f884f5a8b 100644 --- a/lightx2v/models/networks/worldmirror/model.py +++ b/lightx2v/models/networks/worldmirror/model.py @@ -116,22 +116,7 @@ def __init__(self, ln_weight, normalized_shape, eps: float): self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: - w = self.ln_weight._get_actual_weight() - b = self.ln_weight._get_actual_bias() - # Match autocast semantics for a bf16/fp16 activation meeting an - # fp32 LN weight — cast the weight down so the LN runs in the - # activation dtype (what ``nn.LayerNorm`` does under autocast). - if w is not None and x.is_floating_point() and w.dtype != x.dtype and x.dtype in (torch.float16, torch.bfloat16): - w = w.to(x.dtype) - if b is not None: - b = b.to(x.dtype) - return torch.nn.functional.layer_norm( - x, - self.normalized_shape, - w, - b, - self.eps, - ) + return self.ln_weight.apply(x) def _apply(self, fn, *args, **kwargs): # noqa: D401 return self diff --git a/lightx2v/models/networks/worldmirror/weights/transformer_weights.py b/lightx2v/models/networks/worldmirror/weights/transformer_weights.py index 022013d0d..fc2ea9a04 100644 --- a/lightx2v/models/networks/worldmirror/weights/transformer_weights.py +++ b/lightx2v/models/networks/worldmirror/weights/transformer_weights.py @@ -42,6 +42,7 @@ ) _DEFAULT_LN_EPS = 1e-6 +_CAMERA_LN_EPS = 1e-5 # q_norm / k_norm in the ViT blocks operate on head_dim rather than embed_dim # and use a slightly larger epsilon in the source impl (``nn.LayerNorm`` default # = 1e-5) — keeping them separate avoids a silent numerical drift. @@ -62,6 +63,7 @@ def __init__( use_fp32_attn_proj: bool = False, use_fp32_attn_qkv: bool = False, use_fp32_fc1: bool = False, + ln_eps: float = _DEFAULT_LN_EPS, ): super().__init__() self.block_name = block_name @@ -109,7 +111,7 @@ def __init__( LN_WEIGHT_REGISTER[ln_type]( f"{block_name}.norm1.weight", f"{block_name}.norm1.bias", - eps=_DEFAULT_LN_EPS, + eps=ln_eps, ), ) self.add_module( @@ -117,7 +119,7 @@ def __init__( LN_WEIGHT_REGISTER[ln_type]( f"{block_name}.norm2.weight", f"{block_name}.norm2.bias", - eps=_DEFAULT_LN_EPS, + eps=ln_eps, ), ) # --- MLP --- @@ -214,7 +216,8 @@ class WorldMirrorTransformerWeights(WeightModule): ---------- config : dict-like Reads ``dit_quant_scheme`` (default ``"Default"``) and - ``ln_type`` (default ``"torch"``). + ``layer_norm_type`` (default ``"torch"``; legacy + ``ln_norm_type`` is still accepted). depth : int Number of frame / global blocks in ``visual_geometry_transformer``. cam_trunk_depth : int @@ -255,7 +258,7 @@ def __init__( else: mm_type = config.get("dit_quant_scheme", "Default") self._mm_type = mm_type - ln_type = config.get("ln_norm_type", "torch") + ln_type = config.get("layer_norm_type", config.get("ln_norm_type", "torch")) # Block-level fp32 protection: forcibly keep the first # ``fp32_first_n_blocks`` and last ``fp32_last_n_blocks`` of each @@ -360,6 +363,7 @@ def _block_use_fp32(i: int, key: str) -> bool: use_fp32_attn_proj=cam_refine_fp32 or (bool(config.get("use_fp32_attn_proj", False)) if hasattr(config, "get") else False), use_fp32_attn_qkv=cam_refine_fp32 or (bool(config.get("use_fp32_attn_qkv", False)) if hasattr(config, "get") else False), use_fp32_fc1=cam_refine_fp32 or (bool(config.get("use_fp32_fc1", False)) if hasattr(config, "get") else False), + ln_eps=_CAMERA_LN_EPS, ) for i in range(cam_trunk_depth) ), diff --git a/lightx2v/models/networks/z_image/weights/post_weights.py b/lightx2v/models/networks/z_image/weights/post_weights.py index 65f8de9c4..d145f61fb 100755 --- a/lightx2v/models/networks/z_image/weights/post_weights.py +++ b/lightx2v/models/networks/z_image/weights/post_weights.py @@ -18,7 +18,7 @@ def __init__(self, config): "all_final_layer.2-1.adaLN_modulation.1.bias", ), ) - self.add_module("norm_out", LN_WEIGHT_REGISTER["torch"](eps=1e-6)) + self.add_module("norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](eps=1e-6)) self.add_module( "proj_out_linear", diff --git a/lightx2v_platform/ops/norm/enflame_gcu/layer_norm.py b/lightx2v_platform/ops/norm/enflame_gcu/layer_norm.py index 761e5b6df..788ec9025 100644 --- a/lightx2v_platform/ops/norm/enflame_gcu/layer_norm.py +++ b/lightx2v_platform/ops/norm/enflame_gcu/layer_norm.py @@ -6,7 +6,7 @@ @PLATFORM_LAYERNORM_WEIGHT_REGISTER("gcu_layer_norm") class GcuLayerNormWeight(LayerNormWeightTemplate): - def __init__(self, weight_name=None, bias_name=None, create_cuda_buffer=False, create_cpu_buffer=False, lazy_load=False, lazy_load_file=None, is_post_adapter=False, eps=1e-6): + def __init__(self, weight_name=None, bias_name=None, create_cuda_buffer=False, create_cpu_buffer=False, lazy_load=False, lazy_load_file=None, is_post_adapter=False, eps=1e-6, **kwargs): super().__init__(weight_name, bias_name, create_cuda_buffer, create_cpu_buffer, lazy_load, lazy_load_file, is_post_adapter, eps) def apply(self, input_tensor): diff --git a/lightx2v_platform/ops/norm/norm_template.py b/lightx2v_platform/ops/norm/norm_template.py index 3775e5d73..8538a28f4 100644 --- a/lightx2v_platform/ops/norm/norm_template.py +++ b/lightx2v_platform/ops/norm/norm_template.py @@ -253,25 +253,34 @@ def to_cuda(self, non_blocking=False): self.bias = None def to_cpu(self, non_blocking=False): - if hasattr(self, "pin_weight"): + if getattr(self, "pin_weight", None) is not None and getattr(self, "weight", None) is not None: self.weight = self.pin_weight.copy_(self.weight, non_blocking=non_blocking).cpu() - else: + elif getattr(self, "weight", None) is not None: self.weight = self.weight.to("cpu", non_blocking=non_blocking) - if hasattr(self, "pin_bias"): + else: + self.weight = None + if getattr(self, "pin_bias", None) is not None and getattr(self, "bias", None) is not None: self.bias = self.pin_bias.copy_(self.bias, non_blocking=non_blocking).cpu() + elif getattr(self, "bias", None) is not None: + self.bias = self.bias.to("cpu", non_blocking=non_blocking) else: - if self.bias is not None: - self.bias = self.bias.to("cpu", non_blocking=non_blocking) + self.bias = None def state_dict(self, destination=None): if destination is None: destination = {} - destination[self.weight_name] = self.pin_weight if hasattr(self, "pin_weight") else self.weight + if self.weight_name is not None: + destination[self.weight_name] = self.pin_weight if getattr(self, "pin_weight", None) is not None else self.weight if self.bias_name is not None: - destination[self.bias_name] = self.pin_bias if hasattr(self, "pin_bias") else self.bias + destination[self.bias_name] = self.pin_bias if getattr(self, "pin_bias", None) is not None else self.bias return destination def load_state_dict(self, destination, block_index, adapter_block_index=None): + if self.weight_name is None: + self.weight = None + self.bias = None + return + if self.is_post_adapter: assert adapter_block_index is not None weight_name = re.sub(r"\.\d+", lambda m: f".{adapter_block_index}", self.weight_name, count=1) @@ -288,6 +297,11 @@ def load_state_dict(self, destination, block_index, adapter_block_index=None): self.bias = self.bias_cuda_buffer.copy_(destination[bias_name], non_blocking=True) def load_state_dict_from_disk(self, block_index, adapter_block_index=None): + if self.weight_name is None: + self.weight = None + self.bias = None + return + if self.is_post_adapter: self.weight_name = re.sub(r"\.\d+", lambda m: f".{adapter_block_index}", self.weight_name, count=1) self.bias_name = re.sub(r"\.\d+", lambda m: f".{adapter_block_index}", self.bias_name, count=1) if self.bias_name is not None else None diff --git a/test_cases/test_layer_norm_weight_routing.py b/test_cases/test_layer_norm_weight_routing.py new file mode 100644 index 000000000..a0bb1833a --- /dev/null +++ b/test_cases/test_layer_norm_weight_routing.py @@ -0,0 +1,82 @@ +import ast +import importlib.util +import re +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +REPO_ROOT = Path(__file__).resolve().parents[1] +NETWORKS_ROOT = REPO_ROOT / "lightx2v" / "models" / "networks" + + +class LayerNormWeightRoutingTest(unittest.TestCase): + def test_dit_infer_does_not_call_functional_layer_norm_directly(self): + offenders = [] + for path in NETWORKS_ROOT.rglob("*.py"): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "layer_norm": + offenders.append(f"{path.relative_to(REPO_ROOT)}:{node.lineno}") + + self.assertEqual(offenders, [], f"DiT LayerNorm calls must route through registered weights: {offenders}") + + def test_dit_weights_do_not_force_the_torch_layer_norm_registry(self): + hardcoded = [] + pattern = re.compile(r"LN_WEIGHT_REGISTER\[\s*['\"]torch['\"]\s*\]") + for path in NETWORKS_ROOT.rglob("*.py"): + if "weights" not in path.relative_to(NETWORKS_ROOT).parts: + continue + for line_number, line in enumerate(path.read_text().splitlines(), start=1): + if pattern.search(line): + hardcoded.append(f"{path.relative_to(REPO_ROOT)}:{line_number}") + + self.assertEqual(hardcoded, [], f"LayerNorm registry selection must come from config: {hardcoded}") + + def test_platform_layer_norm_supports_no_affine_weight_lifecycle(self): + fake_torch = types.ModuleType("torch") + fake_torch.bfloat16 = object() + fake_torch.float16 = object() + fake_torch.float32 = object() + + fake_safetensors = types.ModuleType("safetensors") + fake_safetensors.safe_open = None + + fake_platform = types.ModuleType("lightx2v_platform") + fake_platform.__path__ = [] + fake_base = types.ModuleType("lightx2v_platform.base") + fake_base.__path__ = [] + fake_global_var = types.ModuleType("lightx2v_platform.base.global_var") + fake_global_var.AI_DEVICE = "cpu" + + stubs = { + "torch": fake_torch, + "safetensors": fake_safetensors, + "lightx2v_platform": fake_platform, + "lightx2v_platform.base": fake_base, + "lightx2v_platform.base.global_var": fake_global_var, + } + module_path = REPO_ROOT / "lightx2v_platform" / "ops" / "norm" / "norm_template.py" + spec = importlib.util.spec_from_file_location("_layer_norm_template_under_test", module_path) + module = importlib.util.module_from_spec(spec) + + with mock.patch.dict(sys.modules, stubs): + spec.loader.exec_module(module) + + class DummyLayerNorm(module.LayerNormWeightTemplate): + def apply(self, input_tensor): + return input_tensor + + layer_norm = DummyLayerNorm(eps=1e-5) + layer_norm.load({}) + layer_norm.to_cpu() + self.assertEqual(layer_norm.state_dict(), {}) + layer_norm.load_state_dict({}, block_index=0) + layer_norm.load_state_dict_from_disk(block_index=0) + self.assertIsNone(layer_norm.weight) + self.assertIsNone(layer_norm.bias) + + +if __name__ == "__main__": + unittest.main() From 35d65ed7e347400e5d394907832f0d302908d457 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yang=20Yong=20=28=E9=9B=8D=E6=B4=8B=29?= Date: Sun, 19 Jul 2026 00:48:55 +0800 Subject: [PATCH 2/2] Delete test_cases/test_layer_norm_weight_routing.py --- test_cases/test_layer_norm_weight_routing.py | 82 -------------------- 1 file changed, 82 deletions(-) delete mode 100644 test_cases/test_layer_norm_weight_routing.py diff --git a/test_cases/test_layer_norm_weight_routing.py b/test_cases/test_layer_norm_weight_routing.py deleted file mode 100644 index a0bb1833a..000000000 --- a/test_cases/test_layer_norm_weight_routing.py +++ /dev/null @@ -1,82 +0,0 @@ -import ast -import importlib.util -import re -import sys -import types -import unittest -from pathlib import Path -from unittest import mock - -REPO_ROOT = Path(__file__).resolve().parents[1] -NETWORKS_ROOT = REPO_ROOT / "lightx2v" / "models" / "networks" - - -class LayerNormWeightRoutingTest(unittest.TestCase): - def test_dit_infer_does_not_call_functional_layer_norm_directly(self): - offenders = [] - for path in NETWORKS_ROOT.rglob("*.py"): - tree = ast.parse(path.read_text(), filename=str(path)) - for node in ast.walk(tree): - if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "layer_norm": - offenders.append(f"{path.relative_to(REPO_ROOT)}:{node.lineno}") - - self.assertEqual(offenders, [], f"DiT LayerNorm calls must route through registered weights: {offenders}") - - def test_dit_weights_do_not_force_the_torch_layer_norm_registry(self): - hardcoded = [] - pattern = re.compile(r"LN_WEIGHT_REGISTER\[\s*['\"]torch['\"]\s*\]") - for path in NETWORKS_ROOT.rglob("*.py"): - if "weights" not in path.relative_to(NETWORKS_ROOT).parts: - continue - for line_number, line in enumerate(path.read_text().splitlines(), start=1): - if pattern.search(line): - hardcoded.append(f"{path.relative_to(REPO_ROOT)}:{line_number}") - - self.assertEqual(hardcoded, [], f"LayerNorm registry selection must come from config: {hardcoded}") - - def test_platform_layer_norm_supports_no_affine_weight_lifecycle(self): - fake_torch = types.ModuleType("torch") - fake_torch.bfloat16 = object() - fake_torch.float16 = object() - fake_torch.float32 = object() - - fake_safetensors = types.ModuleType("safetensors") - fake_safetensors.safe_open = None - - fake_platform = types.ModuleType("lightx2v_platform") - fake_platform.__path__ = [] - fake_base = types.ModuleType("lightx2v_platform.base") - fake_base.__path__ = [] - fake_global_var = types.ModuleType("lightx2v_platform.base.global_var") - fake_global_var.AI_DEVICE = "cpu" - - stubs = { - "torch": fake_torch, - "safetensors": fake_safetensors, - "lightx2v_platform": fake_platform, - "lightx2v_platform.base": fake_base, - "lightx2v_platform.base.global_var": fake_global_var, - } - module_path = REPO_ROOT / "lightx2v_platform" / "ops" / "norm" / "norm_template.py" - spec = importlib.util.spec_from_file_location("_layer_norm_template_under_test", module_path) - module = importlib.util.module_from_spec(spec) - - with mock.patch.dict(sys.modules, stubs): - spec.loader.exec_module(module) - - class DummyLayerNorm(module.LayerNormWeightTemplate): - def apply(self, input_tensor): - return input_tensor - - layer_norm = DummyLayerNorm(eps=1e-5) - layer_norm.load({}) - layer_norm.to_cpu() - self.assertEqual(layer_norm.state_dict(), {}) - layer_norm.load_state_dict({}, block_index=0) - layer_norm.load_state_dict_from_disk(block_index=0) - self.assertIsNone(layer_norm.weight) - self.assertIsNone(layer_norm.bias) - - -if __name__ == "__main__": - unittest.main()