Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model.gguf filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,119 @@
+{
+  "add_classification_head": false,
+  "anyres_pooling_size": 2,
+  "anyres_vit_max_image_size": null,
+  "anyres_vit_two_views": false,
+  "architectures": [
+    "HunYuanMoEV1ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.1,
+  "attention_head_dim": 128,
+  "auto_map": {
+    "AutoConfig": "configuration_hunyuan.HunYuanConfig",
+    "AutoModel": "hunyuan.HunYuanModel",
+    "AutoModelForCausalLM": "hunyuan.HunYuanMoEV1ForCausalLM"
+  },
+  "bos_token_id": 1,
+  "cla_share_factor": 2,
+  "class_num": 0,
+  "dense_list": [
+    4096,
+    0
+  ],
+  "eod_token_id": 127967,
+  "eos_token_id": 127960,
+  "group_limited_greedy": false,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "im_end_id": 6,
+  "im_newline_id": 12,
+  "im_start_id": 5,
+  "image_token_id": 9,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "kv_lora_rank": null,
+  "mask_init_id": 13,
+  "max_position_embeddings": 32768,
+  "mlp_bias": false,
+  "model_type": "hunyuan",
+  "moe_drop_tokens": false,
+  "moe_intermediate_size": [
+    3072,
+    3072,
+    3072,
+    3072
+  ],
+  "moe_layer_num_skipped": 0,
+  "moe_random_routing_dropped_token": false,
+  "moe_topk": [
+    2,
+    2,
+    2,
+    2
+  ],
+  "n_group": null,
+  "norm_topk_prob": true,
+  "norm_type": "rms",
+  "num_attention_heads": 32,
+  "num_experts": 4,
+  "num_hidden_layers": 4,
+  "num_key_value_heads": 8,
+  "num_media_embeds": 257,
+  "num_shared_expert": [
+    1,
+    1,
+    1,
+    1
+  ],
+  "org_vocab_size": 128167,
+  "pad_id": 127961,
+  "pad_token_id": 127961,
+  "pool_type": "last",
+  "position_embedding_xdrope": false,
+  "pretraining_tp": 1,
+  "q_lora_rank": null,
+  "qk_nope_head_dim": null,
+  "qk_rope_head_dim": null,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "alpha": 1000.0,
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 1.0,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "type": "dynamic"
+  },
+  "rope_theta": 10000.0,
+  "routed_scaling_factor": 1.0,
+  "sep_token_id": 127962,
+  "skip_cls_token": false,
+  "text_end_id": 8,
+  "text_start_id": 7,
+  "tie_word_embeddings": true,
+  "topk_group": null,
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.3",
+  "use_cache": true,
+  "use_cla": false,
+  "use_mixed_mlp_moe": true,
+  "use_mla": false,
+  "use_qk_norm": true,
+  "use_rotary_pos_emb": true,
+  "v_head_dim": null,
+  "video_end_id": 11,
+  "video_start_id": 10,
+  "vit_add_patchemb_bias": false,
+  "vit_input_resolution": 224,
+  "vit_mapping_type": "resampler",
+  "vit_norm_type": "fused",
+  "vit_patch": 1,
+  "vit_path": null,
+  "vit_remove_prenorm": false,
+  "vit_token": 64,
+  "vit_type": null,
+  "vit_used_rms_norm": false,
+  "vocab_size": 128167,
+  "xdrope_section": null
+}

configuration_hunyuan.py

@@ -0,0 +1,319 @@
+# coding=utf-8
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved.
+""" HunYuan model configuration"""
+from torch import nn
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from typing import List, Union, Optional
+
+
+logger = logging.get_logger(__name__)
+
+
+class HunYuanConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`HunYuanModel`]. It is used to instantiate an
+    HunYuan model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the HunYuan-7B.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32000):
+            Vocabulary size of the HunYuan model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`HunYuanModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 11008):
+            Dimension of the MLP representations or shared MLP representations.
+        moe_intermediate_size (`int` or `List`, *optional*, defaults to 11008):
+            Dimension of the MLP representations in MoE. Use a list if you want a different size per layer.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*):
+            Padding token id.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            End of stream token id.
+        pretraining_tp (`int`, *optional*, defaults to 1):
+            Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
+            document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
+            necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
+            issue](https://github.com/pytorch/pytorch/issues/76232).
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+            `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
+            these scaling strategies behave:
+            https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
+            experimental feature, subject to breaking API changes in future versions.
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        use_qk_norm (`bool`, *optional*, defaults to `False`):
+            Whether query and key in attention use norm
+        use_cla (`bool`, *optional*, defaults to `False`):
+            Whether to use CLA in attention
+        cla_share_factor (`int`, *optional*, defaults to 1):
+            The share factor of CLA
+        num_experts (`int` or `List`, *optional*, defaults to 1):
+            The number of experts for moe. If it is a list, it will be used as the number of experts for each layer.
+        num_shared_expert (`int` or `List`, *optional*, defaults to 1):
+            The number of shared experts for moe. If it is a list, it will be used as the number of shared experts for each layer.
+        moe_topk (`int` or `List`, *optional*, defaults to 1):
+            The topk value for moe. If it is a list, it will be used as the topk value for each layer.
+        capacity_factor (Not used) (`float` or `List`, *optional*, defaults to 1.0):
+            The capacity factor for moe. If it is a list, it will be used as the capacity factor for each layer.
+        moe_layer_num_skipped (`int`, *optional*, defaults to 0):
+            First moe_layer_num_skipped layers do not use MoE.
+    """
+
+    model_type = "hunyuan"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=290943,
+        org_vocab_size=290943,
+        hidden_size=4096,
+        intermediate_size: int=11008,
+        moe_intermediate_size: Union[int, List]=None,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        attention_head_dim=None,
+        hidden_act="silu",
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        eod_token_id=3,
+        sep_token_id=4,
+        im_start_id=5,
+        im_end_id=6,
+        text_start_id=7,
+        text_end_id=8,
+        image_token_id=9,
+        video_start_id=10,
+        video_end_id=11,
+        im_newline_id=12,
+        mask_init_id=13,
+        pretraining_tp=1,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        mlp_bias=False,
+        attention_dropout=0.0,
+        use_qk_norm=False,
+        use_rotary_pos_emb=True,
+        use_cla=False,
+        cla_share_factor=1,
+        norm_type="hf_rms",
+        num_experts: Union[int, List]=1,
+        use_mixed_mlp_moe=False,
+        num_shared_expert: Union[int, List]=1,
+        moe_topk: Union[int, List]=1,
+        # capacity_factor: Union[int, List]=1.0,
+        moe_drop_tokens=False,
+        moe_random_routing_dropped_token=False,
+        use_mla=False,
+        kv_lora_rank=512,
+        q_lora_rank=1536,
+        qk_rope_head_dim=64,
+        v_head_dim=128,
+        qk_nope_head_dim=128,
+        moe_layer_num_skipped=0,
+        norm_topk_prob=True,
+        routed_scaling_factor=1.0,
+        group_limited_greedy=False,
+        n_group=None,
+        topk_group=None,
+        vit_path=None,
+        num_media_embeds=257,
+        vit_type="AnyResVit",
+        vit_input_resolution=224,
+        vit_token=64,
+        vit_patch=1,
+        vit_mapping_type="simple_conv_mlp",
+        vit_norm_type="fused",
+        vit_used_rms_norm=True,
+        vit_remove_prenorm=True,
+        vit_add_patchemb_bias=True,
+        anyres_vit_max_image_size=2048,
+        anyres_pooling_size=2,
+        anyres_vit_two_views=False,
+        skip_cls_token=False,
+        position_embedding_xdrope=False,
+        xdrope_section=None,
+        add_classification_head=False,
+        class_num=0,
+        pool_type="last",
+        pad_id=-1,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.org_vocab_size = org_vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_experts = num_experts
+        self.use_mixed_mlp_moe = use_mixed_mlp_moe
+        self.num_shared_expert = num_shared_expert
+        self.moe_topk = moe_topk
+        # self.capacity_factor = capacity_factor
+        self.moe_drop_tokens = moe_drop_tokens
+        self.moe_random_routing_dropped_token = moe_random_routing_dropped_token
+
+        if attention_head_dim is not None:
+            self.attention_head_dim = attention_head_dim
+        else:
+            self.attention_head_dim = self.hidden_size // num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.pretraining_tp = pretraining_tp
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        # self._rope_scaling_validation()   # TODO: Need validation?
+        self.attention_bias = attention_bias
+        self.mlp_bias = mlp_bias
+        self.attention_dropout = attention_dropout
+        self.use_qk_norm = use_qk_norm
+        self.use_rotary_pos_emb = use_rotary_pos_emb
+        self.use_cla = use_cla
+        self.cla_share_factor = cla_share_factor
+        self.norm_type = norm_type
+        # MLA args
+        self.use_mla = use_mla
+        self.kv_lora_rank = kv_lora_rank
+        self.q_lora_rank = q_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.v_head_dim = v_head_dim
+
+        # DeepSeek related args
+        self.moe_layer_num_skipped = moe_layer_num_skipped
+        self.norm_topk_prob = norm_topk_prob
+        self.routed_scaling_factor = routed_scaling_factor
+        self.group_limited_greedy = group_limited_greedy
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.add_classification_head = add_classification_head
+        self.class_num = class_num
+        self.pool_type = pool_type
+        self.pad_id = pad_id
+
+        if self.class_num is not None:
+            self.dense_list = [self.hidden_size, self.class_num]
+
+        # Vit args
+        self.vit_path = vit_path
+        self.num_media_embeds = num_media_embeds
+        self.vit_type = vit_type
+        self.vit_input_resolution = vit_input_resolution
+        self.vit_token = vit_token
+        self.vit_patch = vit_patch
+        self.vit_mapping_type = vit_mapping_type
+        self.vit_norm_type = vit_norm_type
+        self.vit_used_rms_norm = vit_used_rms_norm
+        self.vit_remove_prenorm = vit_remove_prenorm
+        self.vit_add_patchemb_bias = vit_add_patchemb_bias
+        self.anyres_vit_max_image_size = anyres_vit_max_image_size
+        self.anyres_pooling_size = anyres_pooling_size
+        self.anyres_vit_two_views = anyres_vit_two_views
+        self.skip_cls_token = skip_cls_token
+        self.position_embedding_xdrope = position_embedding_xdrope
+        self.xdrope_section = xdrope_section
+
+        # token id
+        self.eod_token_id = eod_token_id
+        self.im_start_id = im_start_id
+        self.im_end_id = im_end_id
+        self.text_start_id = text_start_id
+        self.text_end_id = text_end_id
+        self.image_token_id = image_token_id
+        self.video_start_id = video_start_id
+        self.video_end_id = video_end_id
+        self.im_newline_id = im_newline_id
+        self.mask_init_id = mask_init_id
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            sep_token_id=sep_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    def _rope_scaling_validation(self):
+        """
+        Validate the `rope_scaling` configuration.
+        """
+        if self.rope_scaling is None:
+            return
+
+        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
+            raise ValueError(
+                "`rope_scaling` must be a dictionary with with two fields, `type` and `factor` or `type` and `alpha`, "
+                f"got {self.rope_scaling}"
+            )
+        rope_scaling_type = self.rope_scaling.get("type", None)
+        rope_scaling_factor = self.rope_scaling.get("factor", None)
+        rope_scaling_alpha = self.rope_scaling.get("alpha", None)
+        if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
+            raise ValueError(
+                f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
+            )
+        if rope_scaling_factor is None and rope_scaling_alpha is None:
+            raise ValueError("`rope_scaling`'s factor or alpha field must be have one, got both of none")
+        if rope_scaling_factor is not None:
+            if not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
+                raise ValueError(f"`rope_scaling`'s factor field must be a float > 1.0, got {rope_scaling_factor}")
+        if rope_scaling_alpha is not None:
+            if not isinstance(rope_scaling_alpha, float) or rope_scaling_alpha <= 1.0:
+                raise ValueError(f"`rope_scaling`'s alpha field must be a float > 1.0, got {rope_scaling_alpha}")

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 127960,
+  "pad_token_id": 127961,
+  "transformers_version": "4.51.3"
+}

hunyuan.py

@@ -0,0 +1,879 @@
+# coding=utf-8
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved.
+#
+""" PyTorch HunYuan model."""
+
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+    AttentionMaskConverter,
+    _prepare_4d_attention_mask,
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    SequenceClassifierOutputWithPast
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.utils.import_utils import is_torch_fx_available
+from transformers.generation.utils import GenerateOutput
+from .configuration_hunyuan import HunYuanConfig
+from .modeling_hunyuan import HunYuanDecoderLayer, HunYuanRMSNorm
+from .vit_model import NaVitForward, VitForward, Vit 
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+
+# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
+# It means that the function will not be traced through and simply appear as a node in the graph.
+if is_torch_fx_available():
+    if not is_torch_greater_or_equal_than_1_13:
+        import torch.fx
+
+    _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
+
+
+
+_CONFIG_FOR_DOC = "HunYuanConfig"
+
+
+HUNYUAN_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`HunYuanConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare HunYuan Model outputting raw hidden-states without any specific head on top.",
+    HUNYUAN_START_DOCSTRING,
+)
+class HunYuanPreTrainedModel(PreTrainedModel):
+    config_class = HunYuanConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["HunYuanDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+HUNYUAN_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance;
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare HunYuan Model outputting raw hidden-states without any specific head on top.",
+    HUNYUAN_START_DOCSTRING,
+)
+class HunYuanModel(HunYuanPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`HunYuanDecoderLayer`]
+
+    Args:
+        config: HunYuanConfig
+    """
+
+    def __init__(self, config: HunYuanConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.add_classification_head = config.add_classification_head
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [HunYuanDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self._use_sdpa = config._attn_implementation == "sdpa"
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        if not config.add_classification_head:
+            self.norm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.cla = config.use_cla
+        self.cla_share_factor = config.cla_share_factor
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        # if input_ids is not None and inputs_embeds is not None:
+        #     raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        if input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        past_key_values_length = 0
+        if use_cache:
+            use_legacy_cache = not isinstance(past_key_values, Cache)
+            if use_legacy_cache:
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+            past_key_values_length = past_key_values.get_usable_length(seq_length)
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        
+        # Fix lora with gradient checkpointing training
+        if self.training and inputs_embeds.is_leaf:
+            inputs_embeds.requires_grad = True
+
+        if self._use_flash_attention_2:
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self._use_sdpa and not output_attentions:
+            # output_attentions=True can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_key_values_length,
+            )
+        else:
+            # 4d mask is passed through the layers
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            )
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        prev_kv_states = None
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    prev_kv_states,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    kv_states=prev_kv_states
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+            kv_states = layer_outputs[-1]
+
+            if self.cla and layer_idx % self.cla_share_factor == 0:
+                prev_kv_states = kv_states
+        if not self.add_classification_head:
+            hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class HunYuanMoEV1ForCausalLM(HunYuanPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: HunYuanConfig):
+        super().__init__(config)
+        if config.vit_path is not None:
+            if "-tp" in config.vit_type:
+                config.vit_type = config.vit_type.replace("-tp", "")
+            self.vit_type = config.vit_type
+            if self.vit_type not in ['NaVit', 'EvaVit']:
+                if config.vit_mapping_type == 'mlp':
+                    self.vit_linear_encoder = torch.nn.Linear(config.hidden_size, config.hidden_size)
+            self.vit = Vit(config)
+        else:
+            self.vit = None
+        self.config = config
+        self.model = HunYuanModel(config)
+        self.add_classification_head = config.add_classification_head
+        self.pad_id = config.pad_id
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.add_classification_head:
+            self.pool_head = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+            self.pool_head2 = nn.Linear(config.hidden_size, config.class_num, bias=False)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, AutoModelForCausalLM
+
+        >>> model = AutoModelForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+
+        if not self.add_classification_head:
+            if self.config.pretraining_tp > 1:
+                lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
+                logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
+                logits = torch.cat(logits, dim=-1)
+            else:
+                logits = self.lm_head(hidden_states)
+            logits = logits.float()
+        else:
+            logits = hidden_states
+            logits = logits.float()
+            pooled_output = self.pool_head(logits)
+            pooled_output = torch.tanh(pooled_output)
+            pooled_output = self.pool_head2(pooled_output).contiguous()  # bs * class_num
+            if len(pooled_output.shape) < 2:
+                raise ValueError("pooled_output does not have enough dimensions for transpose")
+
+            if self.config.pool_type == "mean":
+                reward = pooled_output.mean(dim=1).squeeze(-1)
+            elif self.config.pool_type == "last":
+                # bs * hidden_size
+                seq_length = (input_ids != self.pad_id).long().sum(dim=1) - 1
+                batch_size = input_ids.size(0)
+                reward = pooled_output[torch.arange(batch_size, device=pooled_output.device), seq_length].squeeze(-1)
+            else:
+                reward = pooled_output[:, 0].squeeze(-1)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.reshape(-1, self.config.vocab_size)
+            shift_labels = shift_labels.reshape(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        output = CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+        if self.add_classification_head:
+            output['reward'] = reward
+
+        return output
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = past_key_values.get_max_cache_shape()
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length):]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1]:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
+
+
+class MultimodelHunYuanForCausalLM(HunYuanMoEV1ForCausalLM):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: HunYuanConfig):
+        super().__init__(config)
+
+    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        imgs: Optional[List[torch.FloatTensor]] = None,
+        imgs_pos: Optional[List[int]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, AutoModelForCausalLM
+
+        >>> model = AutoModelForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        mask_init_id = self.config.mask_init_id
+        pad_id = self.config.pad_token_id
+        eod_id = self.config.eod_token_id
+        image_token_id = self.config.image_token_id
+        im_start_id = self.config.im_start_id
+        im_end_id = self.config.im_end_id
+        video_start_id = self.config.video_start_id
+        video_end_id = self.config.video_end_id
+
+        if self.vit is not None and imgs is not None:
+            encoder_input = self.model.embed_tokens(input_ids)
+            if self.vit_type in ['NaVit', 'EvaVit', 'AnyResVit']:
+                inputs_embeds, input_ids = NaVitForward(input_ids, encoder_input, self.vit, imgs, imgs_pos, self.config.vit_input_resolution, \
+                    im_start_id, im_end_id, image_token_id, self.config.anyres_vit_two_views, self.config.torch_dtype)
+            else:
+                inputs_embeds, input_ids = VitForward(input_ids, encoder_input, self.vit, self.vit_linear_encoder, imgs, imgs_pos, \
+                    self.config.vit_input_resolution, self.config.vit_mapping_type, self.config.vit_patch, self.config.vit_token)
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        if self.config.pretraining_tp > 1:
+            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
+            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
+            logits = torch.cat(logits, dim=-1)
+        else:
+            logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = logits
+            shift_labels = labels
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.reshape(-1, self.config.vocab_size)
+            shift_labels = shift_labels.reshape(-1)
+            shift_tokens = input_ids.reshape(-1)
+            # compute loss
+            mask = (shift_labels < mask_init_id) & (shift_labels != pad_id) & (shift_labels != image_token_id) & (shift_labels != im_start_id) \
+                    & (shift_labels != im_end_id) & (shift_labels != video_start_id) & (shift_labels != video_end_id) & (shift_tokens != pad_id) & (shift_tokens != eod_id)
+            shift_logits = shift_logits[mask, :]
+            shift_labels = shift_labels[mask]
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        imgs = kwargs.pop("imgs", None)
+        imgs_pos = kwargs.pop("imgs_pos", None)
+        inputs = super().prepare_inputs_for_generation(
+            input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs
+        )
+
+        if imgs is not None:
+            inputs['imgs'] = imgs
+        if imgs_pos is not None:
+            inputs['imgs_pos'] = imgs_pos
+        return inputs
+    
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        imgs: Optional[List[torch.FloatTensor]] = None,
+        imgs_pos: Optional[List[int]] = None,
+        **kwargs,
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        if "inputs_embeds" in kwargs:
+            raise NotImplementedError("`inputs_embeds` is not supported")
+        
+        if self.vit is not None:
+            encoder_input = self.model.embed_tokens(inputs)
+            if self.vit_type in ['NaVit', 'EvaVit', 'AnyResVit']:
+                inputs_embeds, input_ids = NaVitForward(inputs, encoder_input, self.vit, imgs, imgs_pos, self.config.vit_input_resolution, \
+                    self.config.im_start_id, self.config.im_end_id, self.config.image_token_id, self.config.anyres_vit_two_views, self.config.torch_dtype)
+            else:
+                inputs_embeds, input_ids = VitForward(inputs, encoder_input, self.vit, self.vit_linear_encoder, imgs, imgs_pos, \
+                    self.config.vit_input_resolution, self.config.vit_mapping_type, self.config.vit_patch, self.config.vit_token)
+
+        return super().generate(
+            inputs=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            eos_token_id=self.config.eod_token_id,
+            **kwargs
+        )
+
+
+@add_start_docstrings(
+    """
+    The HunYuan Model transformer with a sequence classification head on top (linear layer).
+
+    [`HunYuanForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    (e.g. GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """,
+    HUNYUAN_START_DOCSTRING,
+)
+class HunYuanForSequenceClassification(HunYuanPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = HunYuanModel(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
+                    logits.device
+                )
+            else:
+                sequence_lengths = -1
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits.reshape(-1, self.num_labels), labels.reshape(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )

model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0403737be8ae02a16247b7c8e74b06256d82e86393bea8934fd4d4c7ccccf360
+size 4969099160

model-00002-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:999509aca71a5ed94988cbd91253e125e4bcaaffed1b2119e268801ebee176b7
+size 822201896

model.gguf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29ea5aaff3b1c520c12042af37db7ae48cc413a741026d787f8c736e5920199c
+size 2900447776

model.safetensors.index.json

@@ -0,0 +1,105 @@
+{
+  "metadata": {
+    "total_size": 5791289344
+  },
+  "weight_map": {
+    "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.1.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.1.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.1.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.2.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.2.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.2.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.3.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.3.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.experts.3.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.gate.wg.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.shared_mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.shared_mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.shared_mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.key_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.query_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.1.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.1.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.1.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.2.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.2.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.2.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.3.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.3.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.experts.3.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.gate.wg.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.shared_mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.shared_mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.shared_mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.key_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.query_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.1.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.1.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.1.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.2.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.2.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.2.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.3.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.3.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.experts.3.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.gate.wg.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.shared_mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.shared_mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.shared_mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.key_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.query_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.0.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.0.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.0.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.1.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.1.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.1.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.2.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.2.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.2.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.3.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.3.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.experts.3.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.gate.wg.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.shared_mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.shared_mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.mlp.shared_mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.key_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.query_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.norm.weight": "model-00002-of-00002.safetensors"
+  }
+}

modeling_hunyuan.py

@@ -0,0 +1,1756 @@
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved.
+#
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent/Tencent-Hunyuan-Large/blob/main/License.docx
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+""" PyTorch HunYuan model."""
+
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+    AttentionMaskConverter,
+    _prepare_4d_attention_mask,
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    SequenceClassifierOutputWithPast
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.utils.import_utils import is_torch_fx_available
+from .configuration_hunyuan import HunYuanConfig
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+
+# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
+# It means that the function will not be traced through and simply appear as a node in the graph.
+if is_torch_fx_available():
+    if not is_torch_greater_or_equal_than_1_13:
+        import torch.fx
+
+    _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "HunYuanConfig"
+
+
+def topkgating(logits: Tensor, topk: int):
+    logits = logits.float()
+    gates = F.softmax(logits, dim=1)
+    print(f"logits\n", logits)
+    print(f"gates\n", gates)
+    # expert_capacity = topk * gates.shape[0]
+    expert_capacity = max(topk, topk * gates.shape[0] // gates.shape[1])
+    num_experts = int(gates.shape[1])
+    # Top-k router probability and corresponding expert indices for each token.
+    # Shape: [tokens_per_group, num_selected_experts].
+    expert_gate, expert_index = torch.topk(gates, topk)
+    expert_mask = F.one_hot(expert_index, num_experts)
+    # For a given token, determine if it was routed to a given expert.
+    # Shape: [tokens_per_group, num_experts]
+    expert_mask_aux = expert_mask.max(dim=-2)[0]
+    tokens_per_group_and_expert = torch.mean(expert_mask_aux.float(), dim=-2)
+    router_prob_per_group_and_expert = torch.mean(gates.float(), dim=-2)
+    l_aux = num_experts**2 * torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert)
+    print(f"l_aux\n", l_aux)
+
+    gates_s = torch.clamp(
+        torch.matmul(expert_mask.float(), gates.unsqueeze(-1)).sum(dim=1), min=torch.finfo(gates.dtype).eps
+    )
+    print(f"gates_s\n", gates_s)
+    router_probs = gates / gates_s
+    print(f"router_probs\n", router_probs)
+    # Make num_selected_experts the leading axis to ensure that top-1 choices
+    # have priority over top-2 choices, which have priority over top-3 choices,
+    # etc.
+    expert_index = torch.transpose(expert_index, 0, 1)
+    # Shape: [num_selected_experts * tokens_per_group]
+    expert_index = expert_index.reshape(-1)
+
+    # Create mask out of indices.
+    # Shape: [tokens_per_group * num_selected_experts, num_experts].
+    expert_mask = F.one_hot(expert_index, num_experts).to(torch.int32)
+    exp_counts = torch.sum(expert_mask, dim=0).detach()
+
+    # Experts have a fixed capacity that we cannot exceed. A token's priority
+    # within the expert's buffer is given by the masked, cumulative capacity of
+    # its target expert.
+    # Shape: [tokens_per_group * num_selected_experts, num_experts].
+    token_priority = torch.cumsum(expert_mask, dim=0) * expert_mask - 1
+    # Shape: [num_selected_experts, tokens_per_group, num_experts].
+    token_priority = token_priority.reshape((topk, -1, num_experts))
+    # Shape: [tokens_per_group, num_selected_experts, num_experts].
+    token_priority = torch.transpose(token_priority, 0, 1)
+    # For each token, across all selected experts, select the only non-negative
+    # (unmasked) priority. Now, for group G routing to expert E, token T has
+    # non-negative priority (i.e. token_priority[G,T,E] >= 0) if and only if E
+    # is its targeted expert.
+    # Shape: [tokens_per_group, num_experts].
+    token_priority = torch.max(token_priority, dim=1)[0]
+    print(f"token_priority\n", token_priority)
+
+    # Token T can only be routed to expert E if its priority is positive and
+    # less than the expert capacity. One-hot matrix will ignore indices outside
+    # the range [0, expert_capacity).
+    # Shape: [tokens_per_group, num_experts, expert_capacity].
+    valid_mask = torch.logical_and(token_priority >= 0, token_priority < expert_capacity)
+    token_priority = torch.masked_fill(token_priority, ~valid_mask, 0)
+    dispatch_mask = F.one_hot(token_priority, expert_capacity).to(torch.bool)
+    valid_mask = valid_mask.unsqueeze(-1).expand(-1, -1, expert_capacity)
+    dispatch_mask = torch.masked_fill(dispatch_mask, ~valid_mask, 0)
+
+    # The combine array will be used for combining expert outputs, scaled by the
+    # router probabilities. Shape: [num_groups, tokens_per_group, num_experts,
+    # expert_capacity].
+    combine_weights = torch.einsum("...te,...tec->...tec", router_probs, dispatch_mask)
+    exp_counts_capacity = torch.sum(dispatch_mask)
+    exp_capacity_rate = exp_counts_capacity / (logits.shape[0]*topk)
+
+    return [l_aux, exp_capacity_rate], combine_weights, dispatch_mask, exp_counts
+
+
+def top1gating(logits: Tensor, random_routing_dropped_token: bool = False):
+    """Implements Top1Gating on logits."""
+    # everything is in fp32 in this function
+    logits = logits.float()
+    gates = F.softmax(logits, dim=1)
+    capacity = gates.shape[0]
+
+    # Create a mask for 1st's expert per token
+    # noisy gating
+    indices1_s = torch.argmax(gates, dim=1)
+    num_experts = int(gates.shape[1])
+    mask1 = F.one_hot(indices1_s, num_classes=num_experts)
+
+    # gating decisions
+    # exp_counts = torch.sum(mask1, dim=0).detach().to('cpu')
+    exp_counts = torch.sum(mask1, dim=0).detach()
+
+    # Compute l_aux
+    me = torch.mean(gates, dim=0)
+    ce = torch.mean(mask1.float(), dim=0)
+    l_aux = torch.sum(me * ce) * num_experts
+    mask1_rand = mask1
+
+    top_idx = torch.topk(mask1_rand, k=capacity, dim=0)[1]
+
+    new_mask1 = mask1 * torch.zeros_like(mask1).scatter_(0, top_idx, 1)
+    mask1 = new_mask1
+    mask1_bk = mask1
+    if random_routing_dropped_token:
+        not_full = capacity - new_mask1.sum(dim=0)
+        sorted_notfull, indices_notfull = torch.sort(not_full, descending=True)
+        sorted_notfull = sorted_notfull.to(torch.int64)
+        not_full_experts_ids = torch.repeat_interleave(indices_notfull, sorted_notfull)
+        shuffle_not_full_ids = torch.randperm(not_full_experts_ids.shape[0])
+        not_full_experts_ids = not_full_experts_ids[shuffle_not_full_ids]
+        indices1_s_after_drop = torch.argmax(new_mask1, dim=1)
+        # get drop idx
+        drop_mask = 1 - new_mask1.sum(dim=1)
+        drop_mask = drop_mask.bool()
+        drop_idx = drop_mask.nonzero().view(-1)
+        drop_num = drop_mask.sum().to(torch.int64)
+        indices1_s_after_drop.scatter_(0, drop_idx, not_full_experts_ids[:drop_num])
+        nodrop_mask1 = F.one_hot(indices1_s_after_drop, num_classes=num_experts)
+        mask1 = nodrop_mask1
+
+    # Compute locations in capacity buffer
+    locations1 = torch.cumsum(mask1, dim=0) - 1
+
+    # Store the capacity location for each token
+    locations1_s = torch.sum(locations1 * mask1, dim=1)
+
+    # Normalize gate probabilities
+    mask1_float = mask1.float()
+    gates = gates * mask1_float
+
+    locations1_sc = F.one_hot(locations1_s, num_classes=capacity).float()   # one hot to float
+    combine_weights = torch.einsum("se,sc->sec", gates, locations1_sc)
+
+    dispatch_mask = combine_weights.bool()
+
+    exp_counts_capacity = torch.sum(mask1_bk)
+    exp_capacity_rate = exp_counts_capacity / (logits.shape[0])
+    return [l_aux, exp_capacity_rate], combine_weights, dispatch_mask, exp_counts
+
+
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    warnings.warn(
+        "Calling `transformers.models.llama.modeling_llama._prepare_4d_attention_mask` is deprecated and will be "
+        "removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask"
+    )
+    return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
+
+
+def _make_causal_mask(
+    input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+):
+    warnings.warn(
+        "Calling `transformers.models.llama.modeling_llama._make_causal_mask` is deprecated and will be removed in "
+        "v4.37. Use `transformers.models.llama.modeling_llama.AttentionMaskConverter._make_causal_mask"
+    )
+    return AttentionMaskConverter._make_causal_mask(
+        input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length
+    )
+
+
+class HunYuanRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        HunYuanRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+ALL_LAYERNORM_LAYERS.append(HunYuanRMSNorm)
+
+
+class HunYuanRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+        # inv_freq = inv_freq.bfloat16()
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
+
+        self.inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1).float()
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached or self.inv_freq.dtype != torch.float32:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+        return (
+            self.cos_cached[:seq_len].to(dtype=x.dtype),
+            self.sin_cached[:seq_len].to(dtype=x.dtype),
+        )
+
+
+class HunYuanLinearScalingRotaryEmbedding(HunYuanRotaryEmbedding):
+    """HunYuanRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+        t = t / self.scaling_factor
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+class HunYuanDynamicNTKScalingRotaryEmbedding(HunYuanRotaryEmbedding):
+    """
+    HunYuanRotaryEmbedding extended with Dynamic NTK scaling.
+    Credits to the Reddit users /u/bloc97 and /u/emozilla
+    """
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+
+        if seq_len > self.max_position_embeddings:
+            base = self.base * (
+                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
+            ) ** (self.dim / (self.dim - 2))
+            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+class HunYuanDynamicNTKAlphaRotaryEmbedding(HunYuanRotaryEmbedding):
+    """
+    HunYuanRotaryEmbedding extended with Dynamic NTK scaling.
+    Credits to the Reddit users /u/bloc97 and /u/emozilla
+    """
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_alpha=1.0):
+        self.scaling_alpha = scaling_alpha
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        base = self.base * self.scaling_alpha ** (self.dim / (self.dim-2))
+        inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`):
+            The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+            used to pass offsetted position ids when working with a KV-cache.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class HunYuanMLP(nn.Module):
+    def __init__(self, config: HunYuanConfig, layer_idx=None, is_shared_mlp=False):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        if is_shared_mlp:
+            self.intermediate_size = config.intermediate_size * config.num_shared_expert[0]
+        else:
+            self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        if self.config.pretraining_tp > 1:
+            slice = self.intermediate_size // self.config.pretraining_tp
+            gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
+            up_proj_slices = self.up_proj.weight.split(slice, dim=0)
+            down_proj_slices = self.down_proj.weight.split(slice, dim=1)
+
+            gate_proj = torch.cat(
+                [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
+            )
+            up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)
+
+            intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
+            down_proj = [
+                F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)
+            ]
+            down_proj = sum(down_proj)
+        else:
+            down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+        return down_proj
+
+
+class HunYuanTopKGate(nn.Module):
+    def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.moe_topk = config.moe_topk
+        self.drop_tokens = config.moe_drop_tokens
+        self.min_capacity = 8
+        self.random_routing_dropped_token = config.moe_random_routing_dropped_token
+        self.wg = nn.Linear(config.hidden_size, config.num_experts, bias=False, dtype=torch.float32)
+
+    def forward(self, hidden_states):
+        bsz, seq_len, hidden_size = hidden_states.shape
+        hidden_states = hidden_states.reshape(-1, hidden_size)
+        if self.wg.weight.dtype == torch.float32:
+            hidden_states = hidden_states.float()
+        logits = self.wg(hidden_states)
+        if self.moe_topk == 1:
+            gate_output = top1gating(logits, random_routing_dropped_token=self.random_routing_dropped_token)
+        else:
+            gate_output = topkgating(logits, self.moe_topk[0])
+
+        return gate_output
+
+
+class HunYuanMoE(nn.Module):
+    def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.moe_topk = config.moe_topk
+        self.num_experts = config.num_experts
+        if config.use_mixed_mlp_moe:
+            self.shared_mlp = HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=True)
+        self.gate = HunYuanTopKGate(config, layer_idx=layer_idx)
+        self.experts = nn.ModuleList(
+            [HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=False) for _ in range(config.num_experts)]
+        )
+
+    def forward(self, hidden_states):
+        bsz, seq_len, hidden_size = hidden_states.shape
+
+        if self.config.use_mixed_mlp_moe:
+            hidden_states_mlp = self.shared_mlp(hidden_states)
+
+        l_moe, combine_weights, dispatch_mask, exp_counts = self.gate(hidden_states)
+
+        reshaped_input = hidden_states.reshape(-1, hidden_size)
+
+        dispatched_input = torch.einsum("sec,sm->ecm", dispatch_mask.type_as(hidden_states), reshaped_input)
+        print(f"combine_weights\n", combine_weights)
+        print(f"dispatch_mask\n", dispatch_mask)
+        print(f"dispatched_input\n", dispatched_input)
+
+        chunks = dispatched_input.chunk(self.num_experts, dim=0)
+        expert_outputs = []
+        for chunk, expert in zip(chunks, self.experts):
+            expert_outputs.append(expert(chunk))
+
+        expert_output = torch.cat(expert_outputs, dim=0)
+        combined_output = torch.einsum("sec,ecm->sm", combine_weights.type_as(hidden_states), expert_output)
+        combined_output = combined_output.reshape(bsz, seq_len, hidden_size)
+
+        if self.config.use_mixed_mlp_moe:
+            print(f"mlp out\n", hidden_states_mlp)
+            print(f"moe out\n", combined_output)
+            output = hidden_states_mlp + combined_output
+        else:
+            output = combined_output
+
+        return output
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class HunYuanAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        # layer_idx 从 0 开始
+        self.attention_type = 'cross' if config.use_cla and layer_idx % config.cla_share_factor != 0 else 'self'
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+        self.use_qk_norm = config.use_qk_norm
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        if self.attention_type == 'self':
+            self.k_proj = nn.Linear(
+                self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias
+            )
+            self.v_proj = nn.Linear(
+                self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias
+            )
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
+        if self.use_qk_norm:
+            self.query_layernorm = HunYuanRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.key_layernorm = HunYuanRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self._init_rope()
+
+    def _init_rope(self):
+        if self.config.rope_scaling is None:
+            self.rotary_emb = HunYuanRotaryEmbedding(
+                self.head_dim,
+                max_position_embeddings=self.max_position_embeddings,
+                base=self.rope_theta,
+            )
+        else:
+            scaling_type = self.config.rope_scaling["type"]
+            scaling_factor = self.config.rope_scaling["factor"]
+            scaling_alpha = self.config.rope_scaling["alpha"]
+            if scaling_type == "linear":
+                self.rotary_emb = HunYuanLinearScalingRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            elif scaling_type == "dynamic":
+                if scaling_alpha:
+                    self.rotary_emb = HunYuanDynamicNTKAlphaRotaryEmbedding(
+                        self.head_dim,
+                        max_position_embeddings=self.max_position_embeddings,
+                        scaling_alpha=scaling_alpha,
+                        base=self.rope_theta,
+                    )
+                else:
+                    self.rotary_emb = HunYuanDynamicNTKScalingRotaryEmbedding(
+                        self.head_dim,
+                        max_position_embeddings=self.max_position_embeddings,
+                        scaling_factor=scaling_factor,
+                        base=self.rope_theta,
+                    )
+            else:
+                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        kv_states: torch.Tensor = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use "
+                "`attention_mask` instead.`"
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+
+        if self.config.pretraining_tp > 1:
+            query_slices = self.q_proj.weight.split(
+                (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
+            )
+            query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
+            query_states = torch.cat(query_states, dim=-1)
+
+            if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
+                orig_key_states, orig_value_states = kv_states
+                key_states, value_states = kv_states
+            else:
+                key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
+                key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
+                value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
+
+                key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
+                key_states = torch.cat(key_states, dim=-1)
+
+                value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
+                value_states = torch.cat(value_states, dim=-1)
+                orig_key_states, orig_value_states = key_states, value_states
+
+        else:
+            query_states = self.q_proj(hidden_states)
+            if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
+                orig_key_states, orig_value_states = kv_states
+                key_states, value_states = kv_states
+            else:
+                key_states = self.k_proj(hidden_states)
+                value_states = self.v_proj(hidden_states)
+                orig_key_states, orig_value_states = key_states, value_states
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        print(f"query_states\n", query_states)
+        print(f"key_states\n", key_states)
+
+        if self.use_qk_norm:
+            query_states = self.query_layernorm(query_states)
+            key_states = self.key_layernorm(key_states)
+
+        print(f"query_states_normed\n", query_states)
+        print(f"key_states_normed\n", key_states)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        if self.config.pretraining_tp > 1:
+            attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
+            o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
+            attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
+        else:
+            attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value, (orig_key_states, orig_value_states)
+
+
+class HunYuanFlashAttention2(HunYuanAttention):
+    """
+    HunYuan flash attention module. This module inherits from `HunYuanAttention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        kv_states: torch.Tensor = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # HunYuanFlashAttention2 attention does not support output_attentions
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use "
+                "`attention_mask` instead.`"
+            )
+
+            # overwrite attention_mask with padding_mask
+            attention_mask = kwargs.pop("padding_mask")
+
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
+            orig_key_states, orig_value_states = kv_states
+            key_states, value_states = kv_states
+        else:
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+            orig_key_states, orig_value_states = key_states, value_states
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if self.use_qk_norm:
+            query_states = self.query_layernorm(query_states)
+            key_states = self.key_layernorm(key_states)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.attention_dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (HunYuanRMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            # Handle the case where the model is quantized
+            if hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value, (orig_key_states, orig_value_states)
+
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`int`, *optional*):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            causal = self.is_causal and query_length != 1
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+class HunYuanSdpaAttention(HunYuanAttention):
+    """
+    HunYuan attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `HunYuanAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt
+    to SDPA API.
+    """
+
+    # Adapted from HunYuanAttention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        kv_states: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            logger.warning_once(
+                'HunYuanModel is using HunYuanSdpaAttention,'
+                'but `torch.nn.functional.scaled_dot_product_attention`'
+                'does not support `output_attentions=True`. Falling back to the manual attention implementation, '
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. '
+                'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
+            orig_key_states, orig_value_states = kv_states
+            key_states, value_states = kv_states
+        else:
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+            orig_key_states, orig_value_states = key_states, value_states
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if self.use_qk_norm:
+            query_states = self.query_layernorm(query_states)
+            key_states = self.key_layernorm(key_states)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
+        # custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a
+            # causal mask in case q_len == 1.
+            is_causal=self.is_causal and attention_mask is None and q_len > 1,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value, (orig_key_states, orig_value_states)
+
+
+HUNYUAN_ATTENTION_CLASSES = {
+    "eager": HunYuanAttention,
+    "flash_attention_2": HunYuanFlashAttention2,
+    "sdpa": HunYuanSdpaAttention,
+}
+
+
+class HunYuanDecoderLayer(nn.Module):
+    def __init__(self, config: HunYuanConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.layer_idx = layer_idx
+
+        #self.self_attn = HUNYUAN_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
+        self.self_attn = HunYuanAttention(config=config, layer_idx=layer_idx)
+
+        if config.num_experts > 1:
+            self.mlp = HunYuanMoE(config, layer_idx=layer_idx)
+        else:
+            self.mlp = HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=False)
+        self.input_layernorm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        kv_states: Optional[Tuple[torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*):
+                attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+                query_sequence_length, key_sequence_length)` if default attention is used.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+            kv_states (`Tuple(torch.FloatTensor)`, *optional*): Used when CLA is enabled,
+                key and value states from past attention blocks
+        """
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use "
+                "`attention_mask` instead.`"
+            )
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        print(f"Layer {self.layer_idx}\n", hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value, kv_states = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            kv_states=kv_states,
+            **kwargs,
+        )
+        print(f"Layer {self.layer_idx} self_attn output\n", hidden_states)
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        print(f"Layer {self.layer_idx} mlp output\n", hidden_states)
+        hidden_states = residual + hidden_states
+
+        print(f"Layer {self.layer_idx} final output\n", hidden_states)
+
+        exit(0)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        outputs += (kv_states,)
+
+        return outputs
+
+
+HUNYUAN_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`HunYuanConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare HunYuan Model outputting raw hidden-states without any specific head on top.",
+    HUNYUAN_START_DOCSTRING,
+)
+class HunYuanPreTrainedModel(PreTrainedModel):
+    config_class = HunYuanConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["HunYuanDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+HUNYUAN_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance;
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare HunYuan Model outputting raw hidden-states without any specific head on top.",
+    HUNYUAN_START_DOCSTRING,
+)
+class HunYuanModel(HunYuanPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`HunYuanDecoderLayer`]
+
+    Args:
+        config: HunYuanConfig
+    """
+
+    def __init__(self, config: HunYuanConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [HunYuanDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self._use_sdpa = config._attn_implementation == "sdpa"
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self.norm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.cla = config.use_cla
+        self.cla_share_factor = config.cla_share_factor
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        past_key_values_length = 0
+        if use_cache:
+            use_legacy_cache = not isinstance(past_key_values, Cache)
+            if use_legacy_cache:
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+            past_key_values_length = past_key_values.get_usable_length(seq_length)
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        
+        # Fix lora with gradient checkpointing training
+        if self.training and inputs_embeds.is_leaf:
+            inputs_embeds.requires_grad = True
+
+        if self._use_flash_attention_2:
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self._use_sdpa and not output_attentions:
+            # output_attentions=True can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_key_values_length,
+            )
+        else:
+            # 4d mask is passed through the layers
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            )
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        prev_kv_states = None
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    prev_kv_states,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    kv_states=prev_kv_states
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+            kv_states = layer_outputs[-1]
+
+            if self.cla and layer_idx % self.cla_share_factor == 0:
+                prev_kv_states = kv_states
+
+            break # TEST
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class HunYuanMoEV1ForCausalLM(HunYuanPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: HunYuanConfig):
+        super().__init__(config)
+        self.model = HunYuanModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, AutoModelForCausalLM
+
+        >>> model = AutoModelForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        if self.config.pretraining_tp > 1:
+            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
+            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
+            logits = torch.cat(logits, dim=-1)
+        else:
+            logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = past_key_values.get_max_cache_shape()
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length):]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1]:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
+
+
+@add_start_docstrings(
+    """
+    The HunYuan Model transformer with a sequence classification head on top (linear layer).
+
+    [`HunYuanForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    (e.g. GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """,
+    HUNYUAN_START_DOCSTRING,
+)
+class HunYuanForSequenceClassification(HunYuanPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = HunYuanModel(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
+                    logits.device
+                )
+            else:
+                sequence_lengths = -1
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )

test_hunyuan.py

@@ -0,0 +1,40 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+model = AutoModelForCausalLM.from_pretrained("./.", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("./.", trust_remote_code=True)
+
+with torch.no_grad():
+  input_text = "Hi_"
+  inputs = tokenizer(text=input_text, return_tensors="pt")
+  del inputs["token_type_ids"]
+  print(inputs)
+  gen = model.generate(**inputs, max_new_tokens=1, do_sample=False)
+
+  decoded = tokenizer.batch_decode(gen, skip_special_tokens=True)
+  print(decoded)
+
+
+"""
+from hunyuan.configuration_hunyuan import HunYuanConfig
+from hunyuan.modeling_hunyuan import HunYuanMoEV1ForCausalLM
+import torch
+
+config = HunYuanConfig.from_pretrained("./Hunyuan-A13B-Instruct", trust_remote_code=True)
+config.moe_intermediate_size = [3072, 3072]
+config.num_experts = 4
+config.num_shared_expert = [1, 1]
+config.moe_topk = [2, 2]
+config.num_hidden_layers = 4
+
+model = HunYuanMoEV1ForCausalLM(config)
+print(model)
+
+torch.manual_seed(0)
+state_dict = model.state_dict()
+for key in state_dict:
+  state_dict[key].uniform_(-0.2, 0.2)
+model.load_state_dict(state_dict)
+
+model.save_pretrained("./hunyuan-tiny")
+"""

tokenization_hy.py

@@ -0,0 +1,298 @@
+import base64
+import logging
+import os
+import unicodedata
+from typing import Collection, Dict, List, Set, Tuple, Union
+
+import tiktoken
+from transformers import PreTrainedTokenizer, AddedToken
+
+logger = logging.getLogger(__name__)
+
+
+VOCAB_FILES_NAMES = {"vocab_file": "hy.tiktoken"}
+
+PAT_STR = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+# PAT_STR = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+ENDOFTEXT = "<|endoftext|>"
+STARTOFTEXT = "<|startoftext|>"
+BOSTOKEN = "<|bos|>"
+EOSTOKEN = "<|eos|>"
+PADTOKEN = "<|pad|>"
+
+# as the default behavior is changed to allow special tokens in
+# regular texts, the surface forms of special tokens need to be
+# as different as possible to minimize the impact
+EXTRAS = tuple((f"<|extra_{i}|>" for i in range(205)))
+# changed to use actual index to avoid misconfiguration with vocabulary expansion
+
+
+SPECIAL_START_ID = 127957
+
+def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
+    # with open(tiktoken_bpe_file, "rb", encoding="utf-8") as f:
+    #     contents = f.read()
+    dic = {}
+    rank = 0
+    for line in open(tiktoken_bpe_file, "rb"):
+        if line:
+            token, _ = line.split()
+            if base64.b64decode(token) in dic:
+                continue
+            dic[base64.b64decode(token)] = int(rank)
+            rank += 1
+    global SPECIAL_START_ID
+    SPECIAL_START_ID=rank
+    return dic
+
+# NOTE: Please use the code line to check `SPECIAL_START_ID` right, this will affect the SPECIAL_START_ID
+# _load_tiktoken_bpe('/apdcephfs/share_1502809/shaneshu/tokenizer_exp/other_tokenizer_vocab/hy/' + VOCAB_FILES_NAMES['vocab_file'])
+# print(SPECIAL_START_ID)
+
+SPECIAL_TOKENS = tuple(
+    enumerate(
+        (
+            (
+                ENDOFTEXT,
+                STARTOFTEXT,
+                BOSTOKEN,
+                EOSTOKEN,
+                PADTOKEN,
+            )
+            + EXTRAS
+        ),
+        start=SPECIAL_START_ID,
+    )
+)
+# NOTE: Unused Token ID starts from 127962
+SPECIAL_TOKENS_SET = set(t for i, t in SPECIAL_TOKENS)
+
+class HYTokenizer(PreTrainedTokenizer):
+    """hunyuan tokenizer."""
+
+    vocab_files_names = VOCAB_FILES_NAMES
+
+    def __init__(
+        self,
+        vocab_file,
+        errors="replace",
+        extra_vocab_file=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # how to handle errors in decoding UTF-8 byte sequences
+        # use ignore if you are in streaming inference
+        self.errors = errors
+
+        self.mergeable_ranks = _load_tiktoken_bpe(vocab_file)  # type: Dict[bytes, int]
+        self.special_tokens = {
+            token: index
+            for index, token in SPECIAL_TOKENS
+        }
+
+        # try load extra vocab from file
+        if extra_vocab_file is not None:
+            used_ids = set(self.mergeable_ranks.values()) | set(self.special_tokens.values())
+            extra_mergeable_ranks = _load_tiktoken_bpe(extra_vocab_file)
+            for token, index in extra_mergeable_ranks.items():
+                if token in self.mergeable_ranks:
+                    logger.info(f"extra token {token} exists, skipping")
+                    continue
+                if index in used_ids:
+                    logger.info(f'the index {index} for extra token {token} exists, skipping')
+                    continue
+                self.mergeable_ranks[token] = index
+            # the index may be sparse after this, but don't worry tiktoken.Encoding will handle this
+
+        enc = tiktoken.Encoding(
+            "HunYuan",
+            pat_str=PAT_STR,
+            mergeable_ranks=self.mergeable_ranks,
+            special_tokens=self.special_tokens,
+        )
+        assert (
+            len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab
+        ), f"{len(self.mergeable_ranks)} + {len(self.special_tokens)} != {enc.n_vocab} in encoding"
+
+        self.decoder = {
+            v: k for k, v in self.mergeable_ranks.items()
+        }  # type: dict[int, bytes|str]
+        self.decoder.update({v: k for k, v in self.special_tokens.items()})
+
+        self.tokenizer = enc  # type: tiktoken.Encoding
+
+        self.eod_id = self.tokenizer.eot_token
+        self.bod_id = self.special_tokens[STARTOFTEXT]
+        self.bos_id = self.special_tokens[BOSTOKEN]
+        self.eos_id = self.special_tokens[EOSTOKEN]
+        self.pad_id = self.special_tokens[PADTOKEN]
+
+    def __getstate__(self):
+        # for pickle lovers
+        state = self.__dict__.copy()
+        del state["tokenizer"]
+        return state
+
+    def __setstate__(self, state):
+        # tokenizer is not python native; don't pass it; rebuild it
+        self.__dict__.update(state)
+        enc = tiktoken.Encoding(
+            "HunYuan",
+            pat_str=PAT_STR,
+            mergeable_ranks=self.mergeable_ranks,
+            special_tokens=self.special_tokens,
+        )
+        self.tokenizer = enc
+
+    def __len__(self) -> int:
+        return self.tokenizer.n_vocab
+
+    def get_vocab(self) -> Dict[bytes, int]:
+        return self.mergeable_ranks
+
+    def convert_tokens_to_ids(
+        self, tokens: Union[bytes, str, List[Union[bytes, str]]]
+    ) -> List[int]:
+        ids = []
+        if isinstance(tokens, (str, bytes)):
+            if tokens in self.special_tokens:
+                return self.special_tokens[tokens]
+            else:
+                return self.mergeable_ranks.get(tokens)
+        for token in tokens:
+            if token in self.special_tokens:
+                ids.append(self.special_tokens[token])
+            else:
+                ids.append(self.mergeable_ranks.get(token))
+        return ids
+
+    def _add_tokens(
+        self,
+        new_tokens: Union[List[str], List[AddedToken]],
+        special_tokens: bool = False,
+    ) -> int:
+        if not special_tokens and new_tokens:
+            raise ValueError("Adding regular tokens is not supported")
+        for token in new_tokens:
+            surface_form = token.content if isinstance(token, AddedToken) else token
+            if surface_form not in SPECIAL_TOKENS_SET:
+                raise ValueError("Adding unknown special tokens is not supported")
+        return 0
+
+    def save_vocabulary(self, save_directory: str, **kwargs) -> Tuple[str]:
+        """
+        Save only the vocabulary of the tokenizer (vocabulary).
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        file_path = os.path.join(save_directory, "hunyuan.tiktoken")
+        with open(file_path, "w", encoding="utf-8") as w:
+            for k, v in self.mergeable_ranks.items():
+                line = base64.b64encode(k).decode("utf-8") + " " + str(v) + "\n"
+                w.write(line)
+        return (file_path,)
+
+    def tokenize(
+        self,
+        text: str,
+        allowed_special: Union[Set, str] = "all",
+        disallowed_special: Union[Collection, str] = (),
+        **kwargs,
+    ) -> List[Union[bytes, str]]:
+        """
+        Converts a string in a sequence of tokens.
+        Args:
+            text (`str`):
+                The sequence to be encoded.
+            allowed_special (`Literal["all"]` or `set`):
+                The surface forms of the tokens to be encoded as special tokens in regular texts.
+                Default to "all".
+            disallowed_special (`Literal["all"]` or `Collection`):
+                The surface forms of the tokens that should not be in regular texts and trigger errors.
+                Default to an empty tuple.
+            kwargs (additional keyword arguments, *optional*):
+                Will be passed to the underlying model specific encode method.
+        Returns:
+            `List[bytes|str]`: The list of tokens.
+        """
+        tokens = []
+        text = unicodedata.normalize("NFC", text)
+
+        # this implementation takes a detour: text -> token id -> token surface forms
+        for t in self.tokenizer.encode(
+            text, allowed_special=allowed_special, disallowed_special=disallowed_special
+        ):
+            tokens.append(self.decoder[t])
+        return tokens
+
+    def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:
+        """
+        Converts a sequence of tokens in a single string.
+        """
+        text = ""
+        temp = b""
+        for t in tokens:
+            if isinstance(t, str):
+                if temp:
+                    text += temp.decode("utf-8", errors=self.errors)
+                    temp = b""
+                text += t
+            elif isinstance(t, bytes):
+                temp += t
+            else:
+                raise TypeError("token should only be of type types or str")
+        if temp:
+            text += temp.decode("utf-8", errors=self.errors)
+        return text
+
+    @property
+    def vocab_size(self):
+        return self.tokenizer.n_vocab
+
+    def _convert_id_to_token(self, index: int) -> Union[bytes, str]:
+        """Converts an id to a token, special tokens included"""
+        if index in self.decoder:
+            return self.decoder[index]
+        raise ValueError("unknown ids")
+
+    def _convert_token_to_id(self, token: Union[bytes, str]) -> int:
+        """Converts a token to an id using the vocab, special tokens included"""
+        if token in self.special_tokens:
+            return self.special_tokens[token]
+        if token in self.mergeable_ranks:
+            return self.mergeable_ranks[token]
+        raise ValueError("unknown token")
+
+    def _tokenize(self, text: str, **kwargs):
+        """
+        Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for word-based
+        vocabulary or sub-words for sub-word-based vocabularies (BPE/SentencePieces/WordPieces).
+        Do NOT take care of added tokens.
+        """
+        raise NotImplementedError
+
+    def _decode(
+        self,
+        token_ids: Union[int, List[int]],
+        skip_special_tokens: bool = False,
+        errors: str = None,
+        **kwargs,
+    ) -> str:
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+        if skip_special_tokens:
+            token_ids = [i for i in token_ids if i < self.eod_id]
+        return self.tokenizer.decode(token_ids, errors=errors or self.errors)
+
+# tests
+if __name__ == "__main__":
+    tokenizer = HYTokenizer.from_pretrained('./hy')
+    text = '你好，世界'
+    tokens = tokenizer.tokenize(text)
+    print(tokens)
+    ids = tokenizer.convert_tokens_to_ids(tokens)
+    print(ids)
+    text2 = tokenizer.convert_tokens_to_string(tokens)
+    print(text2)
+    ids2 = tokenizer.convert_tokens_to_ids(tokens)

tokenizer_config.json

@@ -0,0 +1,18 @@
+{
+  "architectures": [
+    "GPT2LMHeadModel"
+  ],
+  "model_max_length": 1048576,
+  "tokenizer_class": "HYTokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_hy.HYTokenizer",
+      null
+      ]
+  },
+  "eos_token": "<|eos|>",
+  "model_type": "gpt2",
+  "additional_special_tokens": ["<|startoftext|>", "<|extra_0|>", "<|extra_4|>", "<|extra_5|>", "<|eos|>"],
+  "pad_token": "<|pad|>",
+  "chat_template": "{% set loop_messages = messages %}\n{% if tools %}\n    {% set weekday_map = {'Monday': '星期一', 'Tuesday': '星期二', 'Wednesday': '星期三', 'Thursday': '星期四', 'Friday': '星期五', 'Saturday': '星期六', 'Sunday': '星期日'} %}\n    {% set weekday_cn = weekday_map[strftime_now('%A')] %}\n    {% set datetime_str = strftime_now('%Y-%m-%d %H:%M:%S') %}\n    {% set datetime_str = datetime_str + ' ' + weekday_cn %}\n    {% for message in loop_messages %}\n        {% if 'content' in message %}\n            {% set content = message['content'] %}\n        {% else %}\n            {% set content = '' %}\n        {% endif %}\n        {% if loop.index0 == 0 %}\n            {% set content_tmp = '你是一位函数组合专家。你会得到一个问题和一组可能的函数。根据问题，你需要进行一个或多个函数/工具调用以实现目的。\n如果没有一个函数可以使用，请直接使用自然语言回复用户，以助手：开头。\n如果给定的问题缺少函数所需的参数，请使用自然语言进行提问，向用户询问必要信息，以助手：开头。\n如果调用结果已经足够回答用户问题，请对历史结果进行总结，使用自然语言回复用户，以助手：开头。\n你应该只在工具调用部分返回函数调用。如果你决定调用任何函数，你必须将其格式化为<tool_calls>[{\"name\": \"func_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}},...]</tool_calls>。你不应该在回复中包含任何其他文本。以下是你可以调用的函数列表，格式为JSON。\n' %}\n            {% set content_tmp = content_tmp + '\n' + tools | tojson + '\n' %}\n            {% if message['role'] == 'system' %}\n                {% set content_tmp = content_tmp + '\n额外要求：\n' + content + '\n\n如果你决定返回函数调用，请将其格式化为<tool_calls>[{\"name\": \"func_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}},...]</tool_calls>，不得包含其他文本。如果额外要求里有格式要求，请忽略，以此处为准。\n否则，请参考开头说的三种情况，以助手：开头进行回复。\n\n如果额外要求里有时间信息，就以额外要求里的时间为准，否则，参考当前时间：' + datetime_str %}\n                {% set content = '<|startoftext|>' + content_tmp + '<|extra_4|>' %}\n            {% elif message['role'] == 'user' %}\n                {% set content_tmp = content_tmp + '\n如果你决定返回函数调用，请将其格式化为<tool_calls>[{\"name\": \"func_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}},...]</tool_calls>，不得包含其他文本。\n否则，请参考开头说的三种情况，以助手：开头进行回复。\n\n当前时间：' + datetime_str %}\n                {% set content_tmp = '<|startoftext|>' + content_tmp + '<|extra_4|>'%}\n                {% set content = content_tmp + '用户：' + content + '<|extra_0|>' %}\n            {% endif %}\n        {% else %}\n            {% if message['role'] == 'user' %}\n                {% set content = '用户：' + content + '<|extra_0|>' %}\n            {% elif message['role'] == 'assistant' %}\n                {% if 'tool_calls' in message %}\n                    {% set tool_calls = message['tool_calls'] %}\n                    {% set ns = namespace(tool_calls=\"[\") %}\n                    {% for tool_call in tool_calls %}\n                        {% set function = tool_call['function'] %}\n                        {% set name = function['name'] %}\n                        {% set ns.tool_calls = ns.tool_calls + '{\"name\": \"' + name + '\", '%}\n                        {% set arguments = function['arguments'] %}\n                        {% if arguments is not string %}\n                            {% set arguments = arguments | tojson %}\n                        {% endif %}\n                        {% set ns.tool_calls = ns.tool_calls + '\"arguments\": ' + arguments + '}' %}\n                        {% if not loop.last %}\n                            {% set ns.tool_calls = ns.tool_calls + ', '%}\n                        {% endif %}\n                    {% endfor %}\n                    {% set ns.tool_calls = ns.tool_calls + ']' %}\n                    {% set content = content + '<tool_calls>' + ns.tool_calls + '</tool_calls>' %}\n                {% else %}\n                    {% set content = '助手：' + content %}\n                {% endif %}\n                {% set content = content + '<|eos|>' %}\n            {% elif message['role'] == 'tool' %}\n                {% if content is not string %}\n                    {set content = content | tojson }\n                {% endif %}\n                {% set content = '<tool_response>' + content + '</tool_response>' %}\n                {% set content = content + '<|extra_0|>' %}\n            {% endif %}\n        {% endif %}\n    {{- content -}}\n    {% endfor %}\n{% else %}\n    {% set context = {'has_head': true} %}\n    {% for message in loop_messages %}\n        {% if 'content' in message %}\n            {% set content = message['content'] %}\n        {% else %}\n            {% set content = '' %}\n        {% endif %}\n        {% if loop.index0 == 0 %}\n            {% if content == '' %}\n                {% set _ = context.update({'has_head': false}) %}\n            {% elif message['role'] == 'system' %}\n                {% set content = '<|startoftext|>' + content + '<|extra_4|>' %}\n            {% endif %}\n        {% endif %}\n        {% if message['role'] == 'user' %}\n            {% if loop.index0 == 1 and not context.has_head %}\n                {% set content = '<|startoftext|>' + content %}\n            {% endif %}\n            {% if loop.index0 == 1 and context.has_head %}\n                {% set content = content + '<|extra_0|>' %}\n            {% else %}\n                {% set content = '<|startoftext|>' + content + '<|extra_0|>' %}\n            {% endif %}\n        {% elif message['role'] == 'assistant' %}\n            {% set content = content + '<|eos|>' %}\n        {% elif message['role'] == 'tool' %}\n            {% set content = content + '<|extra_0|>' %}\n        {% endif %}\n        {{- content -}}\n    {% endfor %}\n{% endif %}\n{%- if enable_thinking is defined and enable_thinking is false %}\n    {{- '<think>\\n\\n</think>\\n' }}\n{%- endif %}"
+}

vit_model.py

@@ -0,0 +1,1083 @@
+import json
+import types
+import math
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+from typing import List, Tuple, Optional, Union
+from contextlib import contextmanager
+from transformers.modeling_attn_mask_utils import (
+    _prepare_4d_causal_attention_mask_for_sdpa,  
+    _prepare_4d_causal_attention_mask_for_sdpa, 
+    _prepare_4d_causal_attention_mask,
+)
+from transformers.models.clip.configuration_clip import CLIPVisionConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from .modeling_hunyuan import HunYuanDecoderLayer, HunYuanRMSNorm
+from .configuration_hunyuan import HunYuanConfig
+
+
+def NaVitForward(input_ids, encoder_input, vit, image_tensors, images_pos, vit_input_resolution, im_start_id, im_end_id, image_token_id, anyres_vit_two_views, dtype):
+    # input_ids: (B, L)
+    # encoder_input: (L, B, E)
+    # image_tensors [[Tensor],...,[Tensor]]
+    # image_pos [[Tensor],...,[Tensor]]
+    # tokenizer = get_tokenizer()
+    b = len(input_ids)
+    img_embs = None
+    all_nums = sum([len(tensors) for tensors in image_tensors]) if image_tensors else 0
+    if all_nums != 0:
+        img_embs, img_batch_pos = vit(image_tensors)
+    else:
+        # when no input image, initialize a fake tensor
+        pad_nums = 1
+        image_tensors = [[torch.rand(3, vit_input_resolution, vit_input_resolution, dtype=dtype, device=torch.cuda.current_device()) for _ in range(pad_nums)]]
+        img_embs, img_batch_pos = vit(image_tensors)
+
+    encoder_input = encoder_input.clone()
+    if all_nums > 0:
+        assert len(images_pos) == len(img_batch_pos), \
+                (len(images_pos), len(img_batch_pos))
+        start_token_id = im_start_id
+        end_token_id = im_end_id
+        placeholder_id = image_token_id
+        for idx in range(len(images_pos)):
+            assert len(images_pos[idx]) == len(img_batch_pos[idx]), \
+                (len(images_pos[idx]), len(img_batch_pos[idx]))
+            for p_img_pos_in_batch, p_batch_img_pos in zip(img_batch_pos[idx], images_pos[idx]):
+                # the positions to be filled [s_start, s_end)
+                s_idx, s_start, s_end = p_img_pos_in_batch
+                current_embs = img_embs[s_idx, s_start:s_end]
+                im_s, im_e = p_batch_img_pos
+                assert len(current_embs) == im_e - im_s, \
+                        (img_embs.shape, (s_start, s_end, s_idx), current_embs.shape, (im_s, im_e, idx))
+                if not anyres_vit_two_views:
+                    assert input_ids[idx, im_s - 1] == start_token_id, \
+                            input_ids[idx, im_s - 1]
+                    assert input_ids[idx, im_e] == end_token_id, \
+                            input_ids[idx, im_e]
+                assert (input_ids[idx, im_s:im_e] == placeholder_id).all(), \
+                        f'The tokens to be filled are not the placeholder_id {placeholder_id}: {(input_ids[idx, im_s:im_e] == placeholder_id).sum()} vs {im_e - im_s}'
+                encoder_input[idx, im_s:im_e] = current_embs
+    else:
+        # when no input image, to mask vit value
+        vit_mask = torch.zeros([1, img_embs.shape[0]], device=torch.cuda.current_device())
+        current_embs = img_embs[0, :]
+        encoder_input[0, 1:img_embs.shape[0] + 1] = encoder_input[0, 1:img_embs.shape[0] + 1] * (1 - vit_mask) + current_embs * vit_mask
+    return encoder_input, input_ids
+
+
+def VitForward(input_ids, encoder_input, vit, vit_linear_encoder, image_tensors, images_pos, vit_input_resolution, vit_mapping_type, vit_patch, vit_token):
+    vit_patch_mlp = (vit_patch > 1 and vit_mapping_type == 'mlp') or vit_patch == 0
+
+    b = len(input_ids)
+    if images_pos is None:
+        images_pos = torch.ones([len(input_ids), 1, 3])
+        images_pos[:, :, 1] = images_pos[:, :, 1]*(vit_token + 1)
+        images_pos = images_pos.long()
+
+    real_image_nums = []
+    image_tensors = image_tensors.view(b, -1, 3, vit_input_resolution, vit_input_resolution)
+    real_images = []
+
+    all_nums = 0
+    img_index = []
+    for s in range(len(images_pos)):
+        real_image_num = 0
+        for (im_s, im_e,index) in images_pos[s]:
+            if im_s == -1:
+                break
+            real_image_num += 1
+            all_nums += 1
+            img_index.append(index)
+
+        real_image_nums.append(real_image_num)
+        real_images.append(image_tensors[s][:real_image_num])
+
+    if vit_patch == 1:
+        img_index = None
+
+    if all_nums == 0:
+        # when no input image, initialize a fake tensor
+        img_input = torch.rand(b, 3, vit_input_resolution, vit_input_resolution).cuda().type(image_tensors.dtype)
+        img_embs = vit(img_input)
+        img_embs = vit_linear_encoder(img_embs)
+    else:
+        img_input = torch.cat(real_images)
+        img_embs = vit(img_input, img_index = img_index)
+        img_embs = vit_linear_encoder(img_embs)
+
+    encoder_input = encoder_input.clone()
+    start = 0
+    if all_nums > 0:
+        for s, real_image_len in enumerate(real_image_nums):
+            current_embs = img_embs[start:start + real_image_len, :] #[30, 256, 4096]
+            for ss in range(current_embs.shape[0]):
+                im_s, im_e, index = images_pos[s, ss]
+                # 子图特征更少
+                if index > 0 and vit_patch_mlp:
+                    encoder_input[s, im_s:im_e,] = current_embs[ss, :(im_e-im_s)]
+                else:
+                    encoder_input[s, im_s:im_e] = current_embs[ss, :]
+            start = start + real_image_len
+    else:
+        # when no input image, to mask vit value
+        for s in range(b):
+            vit_mask = torch.zeros([vit_token, 1]).cuda()
+            current_embs = img_embs[:, start:start + 1]
+            encoder_input[1:vit_token + 1, s] = encoder_input[1:vit_token + 1, s] * (1 - vit_mask) + current_embs[:, 0, :] * vit_mask
+            start = start + 1
+    return encoder_input, input_ids
+
+
+def group_images_by_max_seq_len(
+    images: List[List[Tensor]], patch_size: int, 
+    max_seq_len: int, adaptor_patch_size: int, 
+    add_cls_token: bool = False) -> List[List[Tensor]]:
+
+    groups = []
+    group = []
+    pos_groups = []
+    seq_len = 0
+    num_images = 0
+    for image_list in images:
+        pos_group = []
+        for image in image_list:
+            num_images += 1
+            assert isinstance(image, Tensor)
+
+            image_dims = image.shape[-2:]
+            ph, pw = map(lambda t: t // patch_size, image_dims)
+
+            image_seq_len = (ph * pw)
+            new_image_seq_len = image_seq_len
+            grouped_len = seq_len + image_seq_len
+            if add_cls_token:
+                new_image_seq_len += 1
+                grouped_len += num_images
+                
+            assert new_image_seq_len <= max_seq_len, f'image with dimensions {image_dims} exceeds maximum sequence length'
+            
+            if grouped_len > max_seq_len:
+                groups.append(group)
+                group = []
+                seq_len = 0
+                num_images = 1
+
+            group.append(image)
+            start = seq_len // (adaptor_patch_size * adaptor_patch_size)
+            end = start + image_seq_len//(adaptor_patch_size * adaptor_patch_size)
+            batch_idx = len(groups)
+            pos_group.append([batch_idx, start, end])
+            seq_len += image_seq_len
+        pos_groups.append(pos_group)
+
+    if len(group) > 0:
+        groups.append(group)
+
+    return groups, pos_groups
+
+
+class AnyResCLIPVisionEmbeddings(nn.Module):
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+
+        self.config = config
+        # self.sparse_attn_mask = args.sparse_attn_mask
+        # self.use_flash_attn = args.use_flash_attn
+        self.embed_dim = config.hidden_size
+        self.image_size = config.max_image_size
+        self.patch_size = config.patch_size
+        self.max_seq_len = config.max_vit_seq_len
+        self.adaptor_patch_size = config.adaptor_patch_size
+        self.anyres_vit_two_views = config.anyres_vit_two_views
+        self.vit_add_patchemb_bias = config.vit_add_patchemb_bias
+        self.vit_remove_prenorm = config.vit_remove_prenorm
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=self.vit_add_patchemb_bias,
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.skip_cls_token = True
+
+        # add interpolate_pos_encoding
+        if self.anyres_vit_two_views:
+            self.num_positions = self.num_patches
+            self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim) * 0.02)
+        else:
+            self.num_positions = self.num_patches + 1
+            self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)))
+            # self.position_ids = torch.arange(self.num_positions).expand((1, -1))
+            self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+
+        if not self.vit_remove_prenorm:
+            self.pre_layernorm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        """
+        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+        resolution images.
+
+        Source:
+        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+        """
+        num_patches = embeddings.shape[1]
+        position_embeddings = self.position_embedding(self.position_ids)
+        patch_pos_embed = position_embeddings[:, 1:]
+        num_positions = position_embeddings.shape[1] - 1
+        if num_patches == num_positions and height == width:
+             return patch_pos_embed
+        # class_pos_embed = position_embeddings[:, 0]
+        dim = embeddings.shape[-1]
+        h0 = height // self.patch_size
+        w0 = width // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        h0, w0 = h0 + 0.1, w0 + 0.1
+        patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
+        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+        raw_type = patch_pos_embed.dtype
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.to(torch.float32, non_blocking=True),
+            scale_factor=(h0 / math.sqrt(num_positions), w0 / math.sqrt(num_positions)),
+            mode="bilinear",
+            align_corners=False,
+        )
+        patch_pos_embed = patch_pos_embed.to(raw_type, non_blocking=True)
+        assert int(h0) == patch_pos_embed.shape[-2] and int(w0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return patch_pos_embed
+
+    def rescale_positional_embedding(self, out_size):
+        h, w = out_size
+        pos_embed_shape = int((self.position_embedding.shape[1]) ** 0.5)
+        if (h, w) == (pos_embed_shape, pos_embed_shape):
+            return self.position_embedding
+        rescaled_positional_embedding = \
+            self.position_embedding.new_zeros(1, h*w, self.position_embedding.shape[2])
+        pe_2d = self.position_embedding[0].T.contiguous().view(1, -1, pos_embed_shape, pos_embed_shape)
+        pe_2d = F.interpolate(pe_2d, out_size, mode='bilinear', align_corners=False).view(-1, h*w)
+        rescaled_positional_embedding[0] = pe_2d.T.contiguous()
+        return rescaled_positional_embedding
+
+    def forward_single(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        if pixel_values.ndim == 3:
+            pixel_values = pixel_values[None]
+        batch_size, num_channels, height, width = pixel_values.shape
+
+        if self.anyres_vit_two_views:
+            # padding
+            pad_h = (self.patch_size - height % self.patch_size) % self.patch_size
+            pad_w = (self.patch_size - width % self.patch_size) % self.patch_size
+            pixel_values = F.pad(pixel_values, (0, pad_w, 0, pad_h))
+
+        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
+        b, c, h, w = patch_embeds.shape
+
+        # (b, hw, c)
+        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+        if self.anyres_vit_two_views:
+            embeddings = patch_embeds + self.rescale_positional_embedding(out_size=(h, w))
+        else:
+            embeddings = patch_embeds + self.interpolate_pos_encoding(patch_embeds, height, width)
+        if not self.vit_remove_prenorm:
+            embeddings = self.pre_layernorm(embeddings)
+        return embeddings, (h, w)
+
+    def forward(self, images: List[List[Tensor]]):
+        '''
+        Input:
+            images: List[List[Tensor]]
+
+        Return:
+            embeddings: Tensor (B, L, E)
+            attn_mask: Tensor (B, L, 2)
+            pos_groups: List[List[(batch_idx, start, end)]]
+        '''
+        batched_images, pos_groups = group_images_by_max_seq_len(
+            images, self.patch_size, self.max_seq_len, self.adaptor_patch_size, add_cls_token=not self.skip_cls_token)
+        max_seq_len = self.max_seq_len
+
+        # batched_images is a list of a list
+        B = len(batched_images)
+        L = max_seq_len
+        E = self.embed_dim
+
+        embeddings = torch.zeros(B, L, E, dtype=self.config.torch_dtype, requires_grad=True).cuda(non_blocking=True)
+        attn_mask = embeddings.new_full((B, 1, L, L), False, dtype=torch.bool)  # True presents compute
+        assert len(images) == len(pos_groups), (len(images), len(pos_groups))
+
+        batch_images = []
+        batch_pos = []
+        for images_i, pos_group in zip(images, pos_groups):
+            assert len(images_i) == len(pos_group), (len(images_i), len(pos_group))
+            for image, pos in zip(images_i, pos_group):    
+                batch_idx, start, end = pos
+                a2 = self.adaptor_patch_size ** 2
+                # recover the real number of the input image tokens
+                start *= a2
+                end *= a2
+                emb, _ = self.forward_single(image)
+                assert emb.ndim == 3, '(B, L, E)'
+                embeddings[batch_idx, start:end] = emb
+                attn_mask[batch_idx, :, start:end, start:end] = True
+        return embeddings, attn_mask, pos_groups
+
+
+class CLIPVisionEmbeddings(nn.Module):
+    def __init__(self, config: CLIPVisionConfig, add_pre_layernorm=False, skip_cls_token=True, vit_patch=1):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.image_size = config.vit_input_resolution
+        self.patch_size = config.patch_size
+
+        self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=False,
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+
+        self.skip_cls_token = skip_cls_token
+
+        self.num_positions = self.num_patches + 1
+
+        self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)))
+        if vit_patch > 1:
+            self.position_embedding = nn.Embedding(self.num_patches * (vit_patch ** 2 + 1) + 1, self.embed_dim)
+        # 0 支持最大16张图，目前写死了，如需其他的需要额外定义参数
+        elif vit_patch == 0:
+            self.position_embedding = nn.Embedding(self.num_patches * (16 ** 2 + 1) + 1, self.embed_dim)
+        else:
+            self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+
+        if add_pre_layernorm:
+            self.pre_layernorm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        else:
+            self.pre_layernorm = None
+
+    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        """
+        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+        resolution images.
+
+        Source:
+        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+        """
+        num_patches = embeddings.shape[1] - 1
+        position_embeddings = self.position_embedding(self.position_ids)
+        num_positions = position_embeddings.shape[1] - 1
+        if num_patches == num_positions and height == width:
+             return position_embeddings
+        class_pos_embed = position_embeddings[:, 0]
+        patch_pos_embed = position_embeddings[:, 1:]
+        dim = embeddings.shape[-1]
+        h0 = height // self.config.patch_size
+        w0 = width // self.config.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        h0, w0 = h0 + 0.1, w0 + 0.1
+        patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
+        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+        raw_type = patch_pos_embed.dtype
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.float(),
+            scale_factor=(h0 / math.sqrt(num_positions), w0 / math.sqrt(num_positions)),
+            mode="bicubic",
+            align_corners=False,
+        )
+        # print(patch_pos_embed.shape)
+        patch_pos_embed = patch_pos_embed.to(raw_type)
+        assert int(h0) == patch_pos_embed.shape[-2] and int(w0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+
+    def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False, img_index=None) -> torch.Tensor:
+        batch_size, num_channels, height, width = pixel_values.shape
+        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
+        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+        if self.skip_cls_token:
+            embeddings = patch_embeds
+            if img_index is None:
+                position_ids = self.position_ids[:,1:]
+                embeddings = embeddings + self.position_embedding(position_ids)
+            else:
+                position_ids = (torch.tensor(img_index).cuda() * (self.num_positions - 1)).unsqueeze(1).repeat(1, self.num_positions - 1) \
+                    + self.position_ids.expand(batch_size, -1)[:, 1:]
+                embeddings = embeddings + self.position_embedding(position_ids)
+        else:
+            class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+            embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+            if interpolate_pos_encoding:
+                embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+            else:
+                if img_index is None:
+                    embeddings = embeddings + self.position_embedding(self.position_ids)
+                else:
+                    position_ids = self.position_ids.expand(batch_size,-1)[:,0].unsqueeze(1)
+                    new_position = (torch.tensor(img_index).cuda() * (self.num_positions -1)).unsqueeze(1).repeat(1,self.num_positions-1) +  self.position_ids.expand(batch_size,-1)[:,1:]
+                    position_ids = torch.cat([position_ids,new_position],dim=1)
+                    embeddings = embeddings + self.position_embedding(position_ids)
+        if self.pre_layernorm is not None:
+            embeddings = self.pre_layernorm(embeddings)
+        return embeddings
+
+
+class NaVitTransformer(nn.Module):
+    def __init__(self, config: HunYuanConfig, vit_config: CLIPVisionConfig):
+        super().__init__()
+        self.config = config
+        self.vit_config = vit_config
+        with self.prepare_args(config, vit_config):
+            self._use_sdpa = config._attn_implementation == "sdpa"
+            self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+            self.layers = nn.ModuleList(
+                [HunYuanDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+            )
+
+    @contextmanager
+    def prepare_args(self, config, vit_config):
+        hidden_act = config.hidden_act
+        hidden_size = config.hidden_size
+        ffn_hidden_size = config.intermediate_size
+        num_attention_heads = config.num_attention_heads
+        num_key_value_heads = config.num_key_value_heads
+        attention_head_dim = config.attention_head_dim
+        use_qk_norm = config.use_qk_norm
+        use_rotary_pos_emb =  config.use_rotary_pos_emb
+        num_hidden_layers = config.num_hidden_layers   
+        rms_norm_eps = config.rms_norm_eps
+        attention_dropout = config.attention_dropout
+        # hidden_dropout = config.hidden_dropout
+        norm_type = config.norm_type
+        attention_bias = config.attention_bias
+        mlp_bias = config.mlp_bias
+        use_mla = config.use_mla
+        num_experts = config.num_experts
+        _attn_implementation = config._attn_implementation
+        
+        config.hidden_act = vit_config.hidden_act
+        config.hidden_size = vit_config.hidden_size
+        config.intermediate_size = vit_config.intermediate_size
+        config.num_attention_heads = vit_config.num_attention_heads
+        config.num_key_value_heads = None
+        config.attention_head_dim = vit_config.hidden_size // vit_config.num_attention_heads
+        config.use_qk_norm = False
+        config.use_rotary_pos_emb = False
+        config.num_hidden_layers = vit_config.num_hidden_layers
+        config.rms_norm_eps = vit_config.layer_norm_eps
+        config.attention_dropout = vit_config.attention_dropout
+        # config.hidden_dropout = vit_config.hidden_dropout
+        config.norm_type = config.vit_norm_type
+        config.attention_bias = True
+        config.mlp_bias = True
+        config.use_mla = False
+        config.num_experts = 1
+        config._attn_implementation = "eager"
+        
+        yield
+        config.hidden_act = hidden_act
+        config.hidden_size = hidden_size
+        config.intermediate_size = ffn_hidden_size
+        config.num_attention_heads = num_attention_heads
+        config.num_key_value_heads = num_key_value_heads
+        config.attention_head_dim = attention_head_dim
+        config.use_qk_norm = use_qk_norm
+        config.use_rotary_pos_emb = use_rotary_pos_emb
+        config.num_hidden_layers = num_hidden_layers
+        config.rms_norm_eps = rms_norm_eps
+        config.attention_dropout = attention_dropout
+        # config.hidden_dropout = hidden_dropout
+        config.attention_bias = attention_bias
+        config.mlp_bias = mlp_bias
+        config.norm_type = norm_type
+        config.use_mla = use_mla
+        config.num_experts = num_experts
+        config._attn_implementation = _attn_implementation
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        
+        hidden_states, attention_mask, img_pos = self.embeddings(pixel_values)  
+        attention_mask = attention_mask.int()
+        batch_size, seq_length, _ = hidden_states.shape
+        past_key_values_length = 0  
+
+        if self._use_flash_attention_2:
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self._use_sdpa:
+            # output_attentions=True can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                (batch_size, seq_length),
+                hidden_states,
+                past_key_values_length,
+            )
+        else:
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask,
+                (batch_size, seq_length),
+                hidden_states,
+                past_key_values_length,
+            )
+
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask
+            )
+            hidden_states = layer_outputs[0]
+
+        return hidden_states, img_pos
+
+
+class AnyResVitTransformer(NaVitTransformer):
+    def __init__(self, config: HunYuanConfig, vit_config: CLIPVisionConfig, anyres_vit_max_image_size):
+        super().__init__(config, vit_config)
+        old_anyres_vit_max_image_size = vit_config.max_image_size
+        anyres_vit_max_image_size = anyres_vit_max_image_size or old_anyres_vit_max_image_size
+        vit_config.max_image_size = anyres_vit_max_image_size
+        vit_config.torch_dtype = config.torch_dtype
+        vit_config.anyres_vit_two_views = config.anyres_vit_two_views
+        vit_config.vit_remove_prenorm = config.vit_remove_prenorm
+        vit_config.vit_add_patchemb_bias = config.vit_add_patchemb_bias
+        self.embeddings = AnyResCLIPVisionEmbeddings(vit_config)
+        vit_config.max_image_size = old_anyres_vit_max_image_size
+
+    def fix_embeddings_fn(self, pixel_values):
+        # (B, L, E)
+        embeddings, hw = self.embeddings.forward_single(pixel_values)
+        embeddings = self.embeddings.pre_layernorm(embeddings)
+        return embeddings
+
+
+class CLIPVisionTransformer(nn.Module):
+    def __init__(self, config: HunYuanConfig, vit_config: CLIPVisionConfig):
+        super().__init__()
+        embed_dim = vit_config.hidden_size
+
+        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=vit_config.layer_norm_eps)
+        self.embeddings = CLIPVisionEmbeddings(vit_config, skip_cls_token=config.skip_cls_token, vit_patch=config.vit_patch)
+
+        with self.prepare_args(config, vit_config):
+            self.layers = nn.ModuleList(
+                [HunYuanDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+            )
+
+    @contextmanager
+    def prepare_args(self, config, vit_config):
+        hidden_act = config.hidden_act
+        hidden_size = config.hidden_size
+        ffn_hidden_size = config.intermediate_size
+        num_attention_heads = config.num_attention_heads
+        num_key_value_heads = config.num_key_value_heads
+        attention_head_dim = config.attention_head_dim
+        use_qk_norm = config.use_qk_norm
+        use_rotary_pos_emb =  config.use_rotary_pos_emb
+        num_hidden_layers = config.num_hidden_layers   
+        rms_norm_eps = config.rms_norm_eps
+        attention_dropout = config.attention_dropout
+        # hidden_dropout = config.hidden_dropout
+        norm_type = config.norm_type
+        attention_bias = config.attention_bias
+        mlp_bias = config.mlp_bias
+        use_mla = config.use_mla
+        num_experts = config.num_experts
+        _attn_implementation = config._attn_implementation
+
+        config.hidden_act = vit_config.hidden_act
+        config.hidden_size = vit_config.hidden_size
+        config.intermediate_size = vit_config.intermediate_size
+        config.num_attention_heads = vit_config.num_attention_heads
+        config.num_key_value_heads = None
+        config.attention_head_dim = vit_config.hidden_size // vit_config.num_attention_heads
+        config.use_qk_norm = False
+        config.use_rotary_pos_emb = False
+        config.num_hidden_layers = vit_config.num_hidden_layers
+        config.rms_norm_eps = vit_config.layer_norm_eps
+        config.attention_dropout = vit_config.attention_dropout
+        # config.hidden_dropout = 0.0
+        config.norm_type = "fused"
+        config.attention_bias = True
+        config.mlp_bias = True
+        config.use_mla = False
+        config.num_experts = 1
+        config._attn_implementation = "eager"
+
+        yield
+
+        config.hidden_act = hidden_act
+        config.hidden_size = hidden_size
+        config.intermediate_size = ffn_hidden_size
+        config.num_attention_heads = num_attention_heads
+        config.num_key_value_heads = num_key_value_heads
+        config.attention_head_dim = attention_head_dim
+        config.use_qk_norm = use_qk_norm
+        config.use_rotary_pos_emb = use_rotary_pos_emb
+        config.num_hidden_layers = num_hidden_layers
+        config.rms_norm_eps = rms_norm_eps
+        config.attention_dropout = attention_dropout
+        # config.hidden_dropout = hidden_dropout
+        config.norm_type = norm_type
+        config.attention_bias = attention_bias
+        config.mlp_bias = mlp_bias
+        config.use_mla = use_mla
+        config.num_experts = num_experts
+        config._attn_implementation = _attn_implementation
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        interpolate_pos_encoding: Optional[bool] = None,
+        img_index=None
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+
+        """
+        hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, img_index=img_index)
+        hidden_states = self.pre_layrnorm(hidden_states)
+        batch = hidden_states.shape[0]
+        seq_len = hidden_states.shape[1]
+        device = hidden_states.device
+        attention_mask = torch.ones(batch, 1, seq_len, seq_len, dtype=torch.float32, device=device)
+
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask
+            )
+            hidden_states = layer_outputs[0]
+
+        return hidden_states
+
+
+class Vit(torch.nn.Module):
+    def __init__(self, config, resampler_token=64, pool_rate=2):
+        super().__init__()
+        self.config = config
+        self.vit_mapping_type = config.vit_mapping_type
+        self.anyres_vit_max_image_size = config.anyres_vit_max_image_size
+        self.skip_cls_token = config.skip_cls_token
+        self.pool_rate = pool_rate
+        self.vit_type = self.config.vit_type
+        self.anyres_vit_two_views = self.config.anyres_vit_two_views
+        if self.vit_type in ['Vit-g', 'Vit-bigG', 'NaVit', 'EvaVit', 'AnyResVit']:
+            self.img_init(resampler_token, config.vit_input_resolution, config.vit_mapping_type, pool_rate)
+        else:
+            raise NotImplementedError(f"unsupported vit type: {self.vit_type}")
+    
+    def img_init(self, resampler_token=64, vit_input_resolution=224, vit_mapping_type='resampler', pool_rate=2):
+        if self.vit_type == 'AnyResVit':
+            vit_config = json.load(open(f"{self.config.vit_path}/config.json"))
+            self.vit_config = types.SimpleNamespace(**vit_config["vision_config"])
+            self.vit_config.image_size = vit_input_resolution
+            self.vit = AnyResVitTransformer(self.config, self.vit_config, self.anyres_vit_max_image_size)
+        elif self.vit_type == 'Vit-g':
+            vit_config = json.load(open(f"{self.config.vit_path}/config.json"))
+            self.vit_config = types.SimpleNamespace(**{**vit_config["vision_config_dict"],**vit_config["vision_config"]})
+            self.vit_config.vit_input_resolution = vit_input_resolution
+            self.vit = CLIPVisionTransformer(self.config, self.vit_config)  
+        else:
+           assert False, "other vit_types are not supported"
+
+        if self.vit_mapping_type == 'simple_conv_mlp':
+            self.perceive = SimpleConvMlp(self.vit_config.hidden_size, self.config.hidden_size, self.config.anyres_pooling_size, \
+                self.config.vit_used_rms_norm, self.config.rms_norm_eps, poolmlp=False, twoview=True)
+        elif self.vit_mapping_type == 'oryx_mlp':
+            self.perceive = OryxMLPv2(self.vit_config.hidden_size, self.config.hidden_size, twoview=True, use_pe=False)
+        elif self.vit_mapping_type == 'mlp':
+            self.mlp_depth = 2
+            # one mlp layer already in gpt_model.py
+            mlp_hidden_size = self.vit_config.hidden_size
+            if self.vit_type in ['NaVit', 'EvaVit']:
+                mlp_hidden_size *= self.vit_config.adaptor_patch_size **2
+            if self.mlp_depth > 1:
+                mlp_modules = [torch.nn.Linear(mlp_hidden_size, self.config.hidden_size), torch.nn.GELU()]
+                if self.vit_type in ['NaVit', 'EvaVit']:
+                    for _ in range(1, self.mlp_depth):
+                        mlp_modules.append(torch.nn.Linear(self.config.hidden_size, self.config.hidden_size))
+                        mlp_modules.append(torch.nn.GELU())
+                self.perceive = torch.nn.Sequential(*mlp_modules)
+        else:
+           assert False, "other vit_mapping_types are not supported"
+
+        self.vit_patch_mlp = (self.config.vit_patch > 1 and self.vit_mapping_type == 'mlp') or self.config.vit_patch == 0     
+        for name, param in self.named_parameters():
+            setattr(param, "is_vit_param", True)
+
+    def forward(self, images, img_index=None):
+        if self.vit_type in ['AnyResVit']:
+            dtype = self.config.torch_dtype
+            device = torch.cuda.current_device()
+
+            images_size = []
+            for i in range(len(images)):
+                images_size.append([])
+                for j in range(len(images[i])): 
+                    images_size[i].append((images[i][j].size()[1] // self.vit_config.patch_size, images[i][j].size()[2] // self.vit_config.patch_size))
+        
+            images_feats, img_batch_pos = self.vit(pixel_values=images)
+            a2 = self.vit_config.adaptor_patch_size ** 2
+
+            if self.anyres_vit_two_views:
+                step = 2
+            else:
+                step = 1
+            perceive_fn = lambda x, img_size, is_video: self.perceive(x, img_size, is_video=is_video)     
+            images_list = []
+            images_fix_i = 0
+            num_img_batch_pos = len(img_batch_pos)
+            for i in range(num_img_batch_pos): # batch_id
+                for j in range(0, len(img_batch_pos[i]), step): 
+                    if self.anyres_vit_two_views:
+                        lower_idx, lower_begin, lower_end = img_batch_pos[i][j]
+                        lower_begin = lower_begin * a2
+                        lower_end = lower_end * a2
+                        higher_idx, higher_begin, higher_end = img_batch_pos[i][j + 1]
+                        higher_begin = higher_begin * a2
+                        higher_end = higher_end * a2
+                        lower_res_feat = images_feats[lower_idx, lower_begin:lower_end].unsqueeze(0)
+                        higher_res_feat = images_feats[higher_idx, higher_begin:higher_end].unsqueeze(0)
+                        lower_images_size = images_size[i][j]
+                        higher_images_size = images_size[i][j + 1]
+                        images_list.append(self.perceive(lower_res_feat, lower_images_size, higher_res_feat, higher_images_size))
+                    else:
+                        idx, begin, end = img_batch_pos[i][j]
+                        begin = begin * a2
+                        end = end * a2
+                        is_video = hasattr(images[i][j],'_is_video') and images[i][j]._is_video 
+                        images_list.append(perceive_fn(images_feats[idx, begin:end].unsqueeze(0), images_size[i][j], is_video=is_video))
+
+            images = torch.cat(images_list, dim=1)
+
+            new_batch_pos = []
+            k = 0; cur_len = 0
+            for i in range(len(images_size)):
+                new_batch_pos.append([])
+                for j in range(0, len(images_size[i]), step):
+                    new_pos = [0, cur_len, cur_len + images_list[k].size(1)]
+                    cur_len += images_list[k].size(1)
+                    k += 1
+                    new_batch_pos[i].append(new_pos)
+            return images, new_batch_pos
+        elif self.vit_type == 'Vit-g':
+            images = self.vit(pixel_values=images, interpolate_pos_encoding=False, img_index=img_index)
+        else: 
+            assert False, "other vit_types are not supported"
+        
+        if self.vit_mapping_type == 'mlp':
+            if self.vit_type in ['Vit-g'] and not self.skip_cls_token:
+                images = images[:,1:,:]
+            b, v, d = images.shape
+            s = int(math.sqrt(v))
+            images = images.reshape(b, s, s, d)
+
+
+            if self.vit_patch_mlp and img_index is not None:
+                L_tensor = torch.tensor(img_index)
+                device = images.device
+                # 获取子图位置
+                nonzero_indices = torch.nonzero(L_tensor).squeeze().to(device)
+                # 获取主图位置
+                zero_indices = torch.nonzero(L_tensor == 0).squeeze().to(device)
+
+
+                images_nonzero = torch.index_select(images,0, nonzero_indices).to(device)
+                images_zero = torch.index_select(images, 0, zero_indices).to(device)
+
+                # 子图额外多pool一次
+                pool_rate = self.pool_rate * 2
+                images_nonzero = images_nonzero.reshape(-1, s // pool_rate, pool_rate, s // pool_rate, pool_rate, d)
+                images_nonzero = images_nonzero.permute(0, 1, 3, 5, 2, 4).reshape(-1, (s // pool_rate) * (s // pool_rate), d,
+                                                                      pool_rate*pool_rate).mean(-1)
+
+                # 为了组batch折衷方案
+                images_nonzero = F.pad(images_nonzero, (0, 0, 0, (s // self.pool_rate) * (s // self.pool_rate)- (s // pool_rate) * (s // pool_rate)))
+                images_zero = images_zero.reshape(-1, s // self.pool_rate, self.pool_rate, s // self.pool_rate, self.pool_rate, d)
+                images_zero = images_zero.permute(0, 1, 3, 5, 2, 4).reshape(-1, (s // self.pool_rate) * (s // self.pool_rate), d,
+                                                                  self.pool_rate*self.pool_rate).mean(-1)
+                # 组batch
+                images = torch.zeros(b, (s // self.pool_rate) * (s // self.pool_rate), d).to(device).to(images.dtype)
+                images.index_copy_(0, nonzero_indices, images_nonzero)
+                images.index_copy_(0, zero_indices, images_zero)
+
+                if self.mlp_depth >= 2:
+                    images = self.perceive(images)
+            else:
+                if s % self.pool_rate == 0:
+                    images = images.reshape(b, s//self.pool_rate, self.pool_rate, s//self.pool_rate, self.pool_rate, d)
+                    images = images.permute(0, 1, 3, 5, 2, 4).reshape(b, (s//self.pool_rate) * (s//self.pool_rate), d, -1).mean(-1)
+                    if self.mlp_depth >= 2:
+                        images = self.perceive(images)
+                else:
+                    raise ValueError
+        return images
+
+
+class SimpleConvMlp(nn.Module):
+    def __init__(self, in_channels, out_channels, anyres_pooling_size, vit_used_rms_norm, rms_norm_eps, twoview=False, poolmlp=True, cat_extra_token=True):
+        super().__init__()
+
+        embed_std = 1 / math.sqrt(out_channels)
+        if poolmlp:
+            # if args.learnable_mlp_pooling_size is not None:
+            #     in_channels *= args.learnable_mlp_pooling_size ** 2
+            self.proj = nn.Sequential(
+                nn.Linear(in_channels, out_channels),
+                nn.GELU()
+            )
+            self.vit_linear_encoder = nn.Linear(out_channels, out_channels)
+            self.image_newline = nn.Parameter(
+                torch.randn(out_channels) * embed_std
+            )
+        else:
+            self.proj = nn.Sequential(
+                nn.Conv2d(in_channels, in_channels * 2, kernel_size=anyres_pooling_size, stride=anyres_pooling_size),
+                nn.GELU(),
+                nn.Conv2d(in_channels * 2, in_channels * 4, kernel_size=1),
+            )
+            self.mlp = nn.Linear(in_channels * 4, out_channels)
+            self.image_newline = nn.Parameter(
+                torch.randn(in_channels * 4) * embed_std
+            )
+        self.poolmlp = poolmlp
+
+        self.image_begin = nn.Parameter(
+            torch.randn(out_channels) * embed_std
+        )
+        self.image_end = nn.Parameter(
+            torch.randn(out_channels) * embed_std
+        )
+        
+        if twoview:
+            self.image_sep = nn.Parameter(
+                torch.randn(out_channels) * embed_std
+            )
+
+        self.cat_extra_token = cat_extra_token
+        self.use_rms_norm = vit_used_rms_norm
+        if self.use_rms_norm:
+            self.before_rms = HunYuanRMSNorm(in_channels, eps=rms_norm_eps)
+            self.after_rms = HunYuanRMSNorm(out_channels, eps=rms_norm_eps)
+
+    def forward(self, x, size=(16,16), x2=None, size2=(16, 16), is_video=False):
+        return self.single_forward(x=x, size=size, x2=x2, size2=size2, is_video=is_video)
+
+    def single_forward(self, x, size=(16,16), x2=None, size2=(16, 16), is_video=False): 
+        remove_vit_special_tokens = False
+        learnable_mlp_pooling_size = None         
+        if self.use_rms_norm:
+            x = self.before_rms(x)            
+        h, w = size
+        dtype = x.dtype
+        x = x.permute(0, 2, 1).reshape(x.shape[0], -1, h, w)
+        if self.poolmlp:
+            if learnable_mlp_pooling_size is None:
+                x = F.avg_pool2d(x, anyres_pooling_size)
+                x = self.proj(x.permute(0, 2, 3, 1)) # b, h, w, c
+            else:
+                x = x.permute(0, 2, 3, 1)   # b, h, w, c
+                x = x.reshape(x.shape[0], h // learnable_mlp_pooling_size, learnable_mlp_pooling_size, 
+                                    w // learnable_mlp_pooling_size, learnable_mlp_pooling_size, -1)
+                x = x.permute(0, 1, 3, 2, 4, 5).reshape(x.shape[0], h // learnable_mlp_pooling_size, w // learnable_mlp_pooling_size, -1)
+                x = self.proj(x)
+            x = self.vit_linear_encoder(x)
+            b, h, w, c = x.shape
+            if not remove_vit_special_tokens:
+                x = torch.cat([
+                    x,
+                    self.image_newline.reshape(1, 1, 1, c).expand(b, h, 1, c).to(dtype, non_blocking=True)
+                ], dim=2)
+            x = x.reshape(b, -1, c)
+        else:
+            x = self.proj(x) #b,c,h,w
+            if is_video:
+                video_avgpool_size = 2
+                stride = 2
+                x = F.avg_pool2d(x, kernel_size = video_avgpool_size, stride = stride)
+            b, c, h, w = x.shape
+            if not remove_vit_special_tokens:
+                x = torch.cat([
+                    x,
+                    self.image_newline.reshape(1, c, 1, 1).expand(b, c, h, 1).to(dtype, non_blocking=True)
+                ], dim=-1)
+            x = x.reshape(b, c, -1).permute(0, 2, 1)
+            x = self.mlp(x)
+
+
+        if x2 is not None:
+            h2, w2 = size2
+            x2 = x2.permute(0, 2, 1).reshape(x2.shape[0], -1, h2, w2)
+            if self.poolmlp:
+                x2 = F.avg_pool2d(x2, 2)
+                x2 = self.proj(x2.permute(0, 2, 3, 1)) # b, h, w, c
+                x2 = self.vit_linear_encoder(x2)
+                b2, h2, w2, c2 = x2.shape
+                if not remove_vit_special_tokens:
+                    x2 = torch.cat([
+                        x2,
+                        self.image_newline.reshape(1, 1, 1, c2).expand(b2, h2, 1, c2).to(dtype, non_blocking=True)
+                    ], dim=2)
+                x2 = x2.reshape(b2, -1, c2)
+            else:
+                x2 = self.proj(x2)
+                b2, c2, h2, w2 = x2.shape
+                if not remove_vit_special_tokens:
+                    x2 = torch.cat([
+                        x2,
+                        self.image_newline.reshape(1, c2, 1, 1).expand(b2, c2, h2, 1).to(dtype, non_blocking=True)
+                    ], dim=-1)
+                x2 = x2.reshape(b2, c2, -1).permute(0, 2, 1) #b,n,c
+                x2 = self.mlp(x2)
+
+            sep = self.image_sep.reshape(1, 1, -1).expand(b2, 1, x2.shape[-1]).to(dtype, non_blocking=True)
+
+            x = torch.cat([x, sep, x2], dim=1)
+        
+        if self.cat_extra_token:
+            begin = self.image_begin.reshape(1, 1, -1).expand(b, 1, x.shape[-1]).to(dtype, non_blocking=True)
+            end = self.image_end.reshape(1, 1, -1).expand(b, 1, x.shape[-1]).to(dtype, non_blocking=True)
+            x = torch.cat([begin, x, end], dim=1)
+
+        if self.use_rms_norm:
+            return self.after_rms(x)
+        else:
+            return x
+
+
+class NormalizedDwPooler(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        self.predictor = nn.Sequential(
+            nn.Linear(dim*2, dim),
+            nn.GELU(),
+            nn.Linear(dim, dim),
+        )
+    
+    def forward(self, x, forward_type='2x'):
+        B, H, W, C = x.shape
+
+        if forward_type == '2x':
+            new_x = x.reshape(B, H//2, 2, W//2, 2, C).permute(0, 1, 3, 2, 4, 5).reshape(B, H//2, W//2, 4, C)
+            pooled_x = new_x.mean(-2, keepdim=True).expand(-1, -1, -1, 4, -1)
+            fused_x = torch.cat([new_x, pooled_x], dim=-1)
+        elif forward_type == '1x':
+            new_x = x.reshape(B, H, W, 1, C)
+            fused_x = torch.cat([new_x, new_x], dim=-1)
+        elif forward_type == '4x':
+            new_x = x.reshape(B, H//4, 4, W//4, 4, C).permute(0, 1, 3, 2, 4, 5).reshape(B, H//4, W//4, 16, C)
+            pooled_x = new_x.mean(-2, keepdim=True).expand(-1, -1, -1, 16, -1)
+            fused_x = torch.cat([new_x, pooled_x], dim=-1)
+        
+        score = self.predictor(fused_x)
+        normalized_score = F.softmax(score, dim=-2)
+        new_x = (new_x * normalized_score).sum(dim=-2)
+        return new_x
+
+
+class OryxMLPv2(nn.Module):
+    def __init__(self, in_channels, out_channels, twoview=False, use_pe=False):
+        super().__init__()
+        
+        self.proj1 = nn.Linear(in_channels, out_channels)
+        self.proj2 = nn.Linear(out_channels, out_channels)
+        self.act = nn.GELU()
+        self.pooler = NormalizedDwPooler(out_channels)
+        embed_std = 1 / math.sqrt(out_channels)
+
+        self.use_pe = use_pe
+        if not use_pe:
+            self.image_newline = nn.Parameter(
+                torch.randn(out_channels) * embed_std
+            )
+        self.image_begin = nn.Parameter(
+            torch.randn(out_channels) * embed_std
+        )
+        self.image_end = nn.Parameter(
+            torch.randn(out_channels) * embed_std
+        )
+        
+        if twoview:
+            self.image_sep = nn.Parameter(
+                torch.randn(out_channels) * embed_std
+            )
+
+    def forward(self, x, size=(16,16), x2=None, size2=(16, 16), is_video=False):
+        h, w = size
+        dtype = x.dtype
+        x = x.reshape(x.shape[0], h, w, -1)
+        # x = self.pooler(x, forward_type=REGIONAL_POOL)
+        # x = self.proj(x) #b,h,w, c
+        x = self.proj1(x)
+        x = self.pooler(x, forward_type='2x')
+        x = self.act(x)
+        x = self.proj2(x)
+
+
+        b, h, w, c = x.shape
+        if not self.use_pe:
+            x = torch.cat([
+                x,
+                self.image_newline.reshape(1, 1, 1, c).expand(b, h, 1, c).to(dtype)
+            ], dim=2)
+        else:
+            pe_h = torch.arange(h, dtype=torch.long, device=x.device).reshape(1, h, 1, 1).expand(b, h, w, 1).reshape(b, h*w, 1)
+            pe_w = torch.arange(w, dtype=torch.long, device=x.device).reshape(1, 1, w, 1).expand(b, h, w, 1).reshape(b, h*w, 1)
+            pe = torch.cat([pe_h, pe_w], dim=-1)
+
+        x = x.reshape(b, -1, c)
+
+        if x2 is not None:
+            h2, w2 = size2
+            x2 = x2.reshape(x2.shape[0], h2, w2, -1)
+            # x2 = self.pooler(x2, forward_type=REGIONAL_POOL)
+            ## x2 = self.proj(x2) #b,h,w, c
+            x2 = self.proj1(x2)
+            x2 = self.pooler(x2, forward_type='2x')
+            x2 = self.act(x2)
+            x2 = self.proj2(x2)
+
+            b2, h2, w2, c2 = x2.shape
+            if not self.use_pe:
+                x2 = torch.cat([
+                    x2,
+                    self.image_newline.reshape(1, 1, 1, c).expand(b, h2, 1, c).to(dtype)
+                ], dim=2)
+            x2 = x2.reshape(b, -1, c)
+            sep = self.image_sep.reshape(1, 1, -1).expand(b, 1, c2).to(dtype)
+            x = torch.cat([x, sep, x2], dim=1)
+        
+        begin = self.image_begin.reshape(1, 1, -1).expand(b, 1, c).to(dtype)
+        end = self.image_end.reshape(1, 1, -1).expand(b, 1, c).to(dtype)
+        x = torch.cat([begin, x, end], dim=1)
+        # print(x.shape, x2.shape, h, w, h2, w2)
+        # print("vit rank = " + str(torch.distributed.get_rank()) +" x = " + str(x))
+        if self.use_pe:
+            zero_pad = torch.zeros(b, 1, 2, device=x.device, dtype=torch.long)
+            pe = torch.cat([zero_pad, pe, zero_pad], dim=1)
+            assert pe.shape[1] == x.shape[1]
+            return x, pe
+        else:
+            nseq = x.shape[1]
+            fake_pe = torch.zeros(b, nseq, 2, device=x.device, dtype=torch.long)
+            return x #, fake_pe
+