fix mtp_num_layer >= 2 multimodal

[Jintao-Huang]

#29

[gemini-code-assist]

Code Review

This pull request updates the shifts parameter in the roll_tensor function from a hardcoded -1 to -self.layer_number within the mtp_layer.py module. This adjustment ensures that Multi-Token Prediction (MTP) layers correctly align hidden states with the appropriate future token embeddings based on the layer's position. I have no further feedback to provide.

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the Multi-Token Prediction (MTP) implementation by explicitly passing decoder_input through the MTP layers and blocks, replacing the previous method of bundling it with embeddings in a tuple. It also introduces a patch for MultiTokenPredictionBlock.forward to manage hidden state chunking and concatenation across MTP stages. I have no feedback to provide as there were no review comments to evaluate.

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the Multi-Token Prediction (MTP) implementation by passing decoder_input and embedding as separate arguments instead of a tuple, updating gpt_model.py and mtp_layer.py accordingly. It also adds a patch for MultiTokenPredictionBlock.forward in patcher.py to manage multi-stage hidden states. Feedback was provided regarding a potential TypeError in the patched forward method caused by mixing keyword arguments with positional *args, suggesting a more robust positional argument approach.