Currently, the FSDP backend does not support a reference model, which makes it impossible to correctly compute KL loss (e.g., in RLHF or DPO training). Adding support for a reference model will enable correct and efficient KL loss calculation under FSDP.
To ensure reliability, the implementation should also be numerically aligned with the Megatron backend. (Similar to #290).
Currently, the FSDP backend does not support a reference model, which makes it impossible to correctly compute KL loss (e.g., in RLHF or DPO training). Adding support for a reference model will enable correct and efficient KL loss calculation under FSDP.
To ensure reliability, the implementation should also be numerically aligned with the Megatron backend. (Similar to #290).