[megatron] fix text_position_ids (#5783)

Author: Jintao-Huang

swift/megatron/trainers/utils.py

@@ -142,12 +142,9 @@ def get_batch(data_iterator):
     batch = get_batch_on_this_tp_rank(data_iterator)
     args = get_args()
     num_samples = batch.pop('num_samples')
-    position_ids = batch['position_ids']
-    if position_ids.ndim == 3:
-        text_position_ids = position_ids[0]
-        batch['position_ids'] = position_ids[1:]
-    else:
-        text_position_ids = position_ids
+    text_position_ids = batch.pop('text_position_ids', None)
+    if text_position_ids is None:
+        text_position_ids = batch.get('position_ids')
     if args.padding_free and text_position_ids is not None:
         batch['packed_seq_params'] = get_packed_seq_params(text_position_ids)
         batch['packed_seq_params'].num_samples = num_samples