class InternLM2ForCausalLM(nn.Module, DeployModelMixinV1, CudaGraphMixin):    """Rewrote model of InternLM2ForCausalLM."""    packed_modules_mapping = {        'gate_up_proj': [            'w1',            'w3',        ],    }    def __init__(self,                 config: PretrainedConfig,                 ctx_mgr: StepContextManager,                 dtype: torch.dtype = None,                 device: torch.device = None):        super().__init__()        self.config = config        self.ctx_mgr = ctx_mgr        # build Model        self.model = InternLM2Model(config, dtype=dtype, device=device)        # build lm_head        self.output = self.build_lm_head(config.hidden_size, config.vocab_size, bias=False, dtype=dtype, device=device)    def forward(        self,        input_ids: torch.Tensor,        position_ids: torch.Tensor,        past_key_values: List[List[torch.Tensor]],        attn_metadata: Any = None,        inputs_embeds: torch.Tensor = None,        **kwargs,    ):        """Model forward, return logits."""        hidden_states = self.model(            input_ids=input_ids,            position_ids=position_ids,            past_key_values=past_key_values,            attn_metadata=attn_metadata,            inputs_embeds=inputs_embeds,        )        return hidden_states    def get_logits(self, hidden_states: torch.Tensor):        """Compute logits of the model output."""        return self.output(hidden_states)