[megatron] fix pp4 (#4516)

Author: Jintao-Huang

swift/megatron/train/utils.py

@@ -205,7 +205,7 @@ def get_batch(data_iterator):
 
     # TODO: this is pretty hacky, find a better way
     if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()):
-        return None, None, None, None, None
+        return {key: None for key in ['input_ids', 'attention_mask', 'position_ids']}
 
     # get batches based on the TP rank you are on
     batch = get_batch_on_this_tp_rank(data_iterator)