checkpoint inconsistent between original and diffusers

Inizia eseguendo il codice di confronto su MiniMax-H3/FL2VA/transformer e MiniMax-H3/transformer, concentrandoti sui confronti di qkv_proj.weight nel token refiner e nei transformer blocks. Traccia come vengono prodotti o caricati i checkpoint original e diffusers; il lavoro è completato quando la discrepanza di qkv è stata spiegata e corretta, in modo che il checkpoint original produca gli stessi risultati del checkpoint diffusers.

Thanks for opening source such great model! I have found the checkpoint under MiniMax-H3/FL2VA/transformer is inconsistent with the checkpoint under MiniMax-H3/transformer. The qkv linear weight is different and the others are same. Here is my comparison code

import torch
from safetensors.torch import load_file


if __name__ == '__main__':

    state_dict = {}
    state_dict2 = {}
    for i in range(1,14):
        state_dict.update(load_file('MiniMax-H3/FL2VA/transformer/model-%05d-of-00013.safetensors'%i))
    

    for i in range(1,15):
        state_dict2.update(load_file('MiniMax-H3/transformer/diffusion_pytorch_model-%05d-of-00014.safetensors'%i))
    
    print((state_dict['video_patch_proj.weight']-state_dict2['proj_in.weight']).abs().sum())
    print((state_dict['video_patch_proj.bias']-state_dict2['proj_in.bias']).abs().sum())

    print((state_dict['audio_patch_proj.weight']-state_dict2['audio_proj_in.weight']).abs().sum())
    print((state_dict['audio_patch_proj.bias']-state_dict2['audio_proj_in.bias']).abs().sum())

    print((state_dict['condition_proj.weight']-state_dict2['context_embedder.weight']).abs().sum())
    print((state_dict['condition_proj.bias']-state_dict2['context_embedder.bias']).abs().sum())
    
    print((state_dict['time_embedder.proj_in.weight']-state_dict2['time_embedder.linear_1.weight']).abs().sum())
    print((state_dict['time_embedder.proj_in.bias']-state_dict2['time_embedder.linear_1.bias']).abs().sum())

    print((state_dict['time_embedder.proj_out.weight']-state_dict2['time_embedder.linear_2.weight']).abs().sum())
    print((state_dict['time_embedder.proj_out.bias']-state_dict2['time_embedder.linear_2.bias']).abs().sum())

    for i in range(2):
        print((state_dict['token_refiner.blocks.%d.norm1.weight'%i]-state_dict2['token_refiner.refiner_blocks.%d.norm1.weight'%i]).abs().sum())
        print((state_dict['token_refiner.blocks.%d.norm2.weight'%i]-state_dict2['token_refiner.refiner_blocks.%d.norm2.weight'%i]).abs().sum())

        print((state_dict['token_refiner.blocks.%d.attn.q_norm.weight'%i]-state_dict2['token_refiner.refiner_blocks.%d.attn.norm_q.weight'%i]).abs().sum())
        print((state_dict['token_refiner.blocks.%d.attn.k_norm.weight'%i]-state_dict2['token_refiner.refiner_blocks.%d.attn.norm_k.weight'%i]).abs().sum())

        qkv = torch.cat([state_dict2['token_refiner.refiner_blocks.%d.attn.to_q.weight'%i], state_dict2['token_refiner.refiner_blocks.%d.attn.to_k.weight'%i], state_dict2['token_refiner.refiner_blocks.%d.attn.to_v.weight'%i]], dim=0)
        print((state_dict['token_refiner.blocks.%d.attn.qkv_proj.weight'%i].to(torch.float32) - qkv.to(torch.float32)).abs().mean())

        print((state_dict['token_refiner.blocks.%d.attn.out_proj.weight'%i]-state_dict2['token_refiner.refiner_blocks.%d.attn.to_out.0.weight'%i]).abs().sum())

        print((state_dict['token_refiner.blocks.%d.mlp.fc1.weight'%i]-state_dict2['token_refiner.refiner_blocks.%d.ff.net.0.proj.weight'%i]).abs().sum())
        print((state_dict['token_refiner.blocks.%d.mlp.fc2.weight'%i]-state_dict2['token_refiner.refiner_blocks.%d.ff.net.2.weight'%i]).abs().sum())
    
    print((state_dict['token_refiner.final_norm.weight']-state_dict2['token_refiner.final_norm.weight']).abs().sum())

    for i in range(50):
        print((state_dict['blocks.%d.norm1.weight'%i]-state_dict2['transformer_blocks.%d.norm1.weight'%i]).abs().sum())
        print((state_dict['blocks.%d.norm2.weight'%i]-state_dict2['transformer_blocks.%d.norm2.weight'%i]).abs().sum())

        print((state_dict['blocks.%d.adaln_proj.linear.weight'%i]-state_dict2['transformer_blocks.%d.adaln_proj.linear.weight'%i]).abs().sum())
        print((state_dict['blocks.%d.adaln_proj.linear.bias'%i]-state_dict2['transformer_blocks.%d.adaln_proj.linear.bias'%i]).abs().sum())

        print((state_dict['blocks.%d.attn.q_norm.weight'%i]-state_dict2['transformer_blocks.%d.attn.norm_q.weight'%i]).abs().sum())
        print((state_dict['blocks.%d.attn.k_norm.weight'%i]-state_dict2['transformer_blocks.%d.attn.norm_k.weight'%i]).abs().sum())

        qkv = torch.cat([state_dict2['transformer_blocks.%d.attn.to_q.weight'%i], state_dict2['transformer_blocks.%d.attn.to_k.weight'%i], state_dict2['transformer_blocks.%d.attn.to_v.weight'%i]], dim=0)
        print((state_dict['blocks.%d.attn.qkv_proj.weight'%i].to(torch.float32) - qkv.to(torch.float32)).abs().mean())

        print((state_dict['blocks.%d.attn.out_proj.weight'%i]-state_dict2['transformer_blocks.%d.attn.to_out.0.weight'%i]).abs().sum())

        print((state_dict['blocks.%d.mlp.fc1.weight'%i]-state_dict2['transformer_blocks.%d.ff.net.0.proj.weight'%i]).abs().sum())
        print((state_dict['blocks.%d.mlp.fc2.weight'%i]-state_dict2['transformer_blocks.%d.ff.net.2.weight'%i]).abs().sum())
    

    print((state_dict['final_layer.norm.weight']-state_dict2['norm_out.norm.weight']).abs().sum())
    print((state_dict['final_layer.adaln_proj.linear.weight']-state_dict2['norm_out.linear.weight']).abs().sum())
    print((state_dict['final_layer.adaln_proj.linear.bias']-state_dict2['norm_out.linear.bias']).abs().sum())
    print((state_dict['final_layer.video_out.weight']-state_dict2['proj_out.weight']).abs().sum())
    print((state_dict['final_layer.video_out.bias']-state_dict2['proj_out.bias']).abs().sum())
    print((state_dict['final_layer.audio_out.weight']-state_dict2['audio_proj_out.weight']).abs().sum())
    print((state_dict['final_layer.audio_out.bias']-state_dict2['audio_proj_out.bias']).abs().sum())

And I have checked the SHA256 of the downloaded checkpoint is consistent with the checkpoint in huggingface. I can get normal results using the diffusers checkpoint, but the results are incorrect when using the original checkpoint.