Enable naive Pipeline Parallelism training for Gpt neox japanese and san japanese (#22702)

Move labels to same device as logits

src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py

@@ -682,6 +682,9 @@ class GPTNeoXJapaneseForCausalLM(GPTNeoXJapanesePreTrainedModel):
 
         lm_loss = None
         if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(lm_logits.device)
+
             # we are doing next-token prediction; shift prediction scores and input ids by one
             shift_logits = lm_logits[:, :-1, :].contiguous()
             labels = labels[:, 1:].contiguous()

src/transformers/models/gptsan_japanese/modeling_gptsan_japanese.py

@@ -1236,6 +1236,9 @@ class GPTSanJapaneseForConditionalGeneration(GPTSanJapanesePreTrainedModel):
         router_probs = None
         aux_loss = None
         if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(lm_logits.device)
+
             loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
 
             if output_router_logits: