fix

Author: hjh0119

swift/megatron/trainers/gkd_trainer.py

@@ -348,7 +348,9 @@ def _compute_teacher_logits_from_api(self, encoded_batches: List[Dict]) -> None:
         rollout_src = torch.distributed.get_global_rank(rollout_group, 0)
 
         for encoded_batch in encoded_batches:
-            input_ids = encoded_batch['input_ids']
+            opsd_batch = encoded_batch.get('opsd_teacher_batch')
+            source = opsd_batch if opsd_batch is not None else encoded_batch
+            input_ids = source['input_ids']
             device = input_ids.device
 
             if rollout_rank == 0:
@@ -370,7 +372,6 @@ def _compute_teacher_logits_from_api(self, encoded_batches: List[Dict]) -> None:
             encoded_batch['teacher_api_indices'] = teacher_indices
             encoded_batch['teacher_logits'] = None
 
-            opsd_batch = encoded_batch.get('opsd_teacher_batch')
             if opsd_batch is not None:
                 encoded_batch['opsd_teacher_labels'] = opsd_batch.get('labels')
 

swift/rlhf_trainers/gkd_trainer.py

@@ -590,7 +590,9 @@ def _fetch_teacher_logprobs_from_api(self, encoded_inputs: Dict[str, torch.Tenso
         Returns:
             Tuple of (teacher_logprobs, teacher_indices) tensors with shapes [batch, seq_len, topk]
         """
-        input_ids = encoded_inputs['input_ids']
+        opsd_teacher_inputs = encoded_inputs.get('_opsd_teacher_inputs')
+        source = opsd_teacher_inputs if opsd_teacher_inputs is not None else encoded_inputs
+        input_ids = source['input_ids']
         teacher_logprobs, teacher_indices = fetch_teacher_logprobs(
             self.teacher_model_server, input_ids.tolist(), topk=self.gkd_logits_topk)
         return teacher_logprobs.to(input_ids.device), teacher_indices.to(input_ids.device)