[Megatron-SWIFT] Does RLHF support loss_scale?

Checklist / 检查清单

• I have searched existing issues, and this is a new question or discussion topic. / 我已经搜索过现有的 issues，确认这是一个新的问题与讨论。

Question Description / 问题描述

I'm trying to understand the interaction between loss_scale and Megatron's RLHF trainers (DPO/KTO specifically).

I noticed both trainers explicitly pop the loss_scale key from batch data, and the KTO trainer even has a comment indicating "not support loss_scale". Relevant code snippets are as follows:

• megatron/trainers/kto_trainer.py: 122

with self.stimer(bdata=True):
    # not support loss_scale
    data, kl_data = self.get_batch(data_iterator, vp_stage)
timers('batch-generator').stop()
label = data.pop('label')
data.pop('loss_scale', None)
kl_data.pop('loss_scale', None)

• megatron/trainers/dpo_trainer.py: 122

with self.stimer(bdata=True):
    data = self.get_batch(data_iterator, vp_stage)
timers('batch-generator').stop()
data.pop('loss_scale', None)

Follow-up questions

1. Does the loss_scale mechanism natively support Megatron's RLHF trainers (including DPO/KTO)?
2. If supported, what is the specific algorithm-level implementation logic? For example, in the DPO trainer, are both positive and negative samples masked with loss_scale?

Start by reading megatron/trainers/kto_trainer.py around line 122 and megatron/trainers/dpo_trainer.py around line 122, focusing on where loss_scale is removed from the batch data. Trace the RLHF trainer loss calculations to determine whether loss_scale is supported and how positive and negative samples are handled. Done means documenting the native support and the algorithm-level masking behavior for DPO and KTO.