support batchified-training

modeling_chartmoe.py

@@ -16,7 +16,7 @@
 # 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 InternLMXComposer2 model."""
+"""PyTorch ChartMoE model."""
 import copy
 import queue
 import threading
@@ -262,15 +262,30 @@ class ChartMoEForCausalLM(InternLM2PreTrainedModel):
             wrap_target = wrap_target[:, :self.max_length].to(self.device)
             wrap_im_mask = wrap_im_mask[:, :self.max_length].to(self.device)
 
-            wrap_embeds_list.append(wrap_embeds)
-            wrap_atts_list.append(wrap_atts)
-            wrap_target_list.append(wrap_target)
-            wrap_im_mask_list.append(wrap_im_mask)
-
-        wrap_embeds = torch.cat(wrap_embeds_list)
-        wrap_atts = torch.cat(wrap_atts_list)
-        wrap_target = torch.cat(wrap_target_list)
-        wrap_im_mask = torch.cat(wrap_im_mask_list)
+            wrap_embeds_list.append(wrap_embeds.squeeze(dim=0))
+            wrap_atts_list.append(wrap_atts.squeeze(dim=0))
+            wrap_target_list.append(wrap_target.squeeze(dim=0))
+            wrap_im_mask_list.append(wrap_im_mask.squeeze(dim=0))
+
+        # get the embed of padding token
+        pad = torch.ones([1, 1]) * self.tokenizer.pad_token_id
+        pad = pad.long().to(self.device)
+        pad_emb = self.model.tok_embeddings(pad).squeeze(dim=0)
+        
+        # pack and pad
+        longest_len = max(t.size(0) for t in wrap_embeds_list)
+        wrap_embeds = []
+        for wrap_embed in wrap_embeds_list:
+            pad_len = longest_len - wrap_embed.size(0)
+            if pad_len > 0:
+                pad_tensor = pad_emb.repeat(pad_len, 1)
+                wrap_embed = torch.cat([wrap_embed, pad_tensor], dim=0)
+            wrap_embeds.append(wrap_embed)
+        wrap_embeds = torch.stack(wrap_embeds, dim=0)
+        
+        wrap_atts = nn.utils.rnn.pad_sequence(wrap_atts_list, batch_first=True, padding_value=0.)
+        wrap_target = nn.utils.rnn.pad_sequence(wrap_target_list, batch_first=True, padding_value=-100).long()
+        wrap_im_mask = nn.utils.rnn.pad_sequence(wrap_im_mask_list, batch_first=True, padding_value=0.)
         return wrap_embeds, wrap_atts, wrap_target, wrap_im_mask
 
     def mask_human_targets(self, input_ids, pure=False):