GKD trainer + chunked JSD loss + FSDP

Nobody has claimed this yet.

Hello Liger Kernel team,

First of all, thank you for making this project available! I’ve been exploring your codebase and tried to implement GKDTrainer using the chunked_jsd_loss similarly to how ORPOTrainer handles it. I’m now aiming to use Fully Sharded Data Parallel (FSDP) for both the teacher and student models but am unsure of the best way to integrate it.

I would greatly appreciate any guidance you could provide on:

Implementing the chunked JSD loss function for FSDP-enabled training – Are there recommended patterns or helper functions within the codebase that can simplify this process?
Key code structures or APIs in the GKDTrainer – Which parts of GKDTrainer might need modification or extension to properly handle chunked JSD loss under FSDP?
Best practices or potential pitfalls – Have you encountered any common issues or gotchas when combining chunked losses with FSDP that I should be aware of?
Code snippets or references – If you have any example snippets, documentation references, or design patterns that illustrate how to properly handle teacher and student models together under FSDP, that would be incredibly helpful.
Thank you in advance for your time and assistance! Any insights, tips, or examples you can share will help me get up and running much more quickly.

Additional Context:

I’m currently referencing the ORPOTrainer sample but see that it doesn’t fully address the GKD use case.

Start by comparing the GKDTrainer entry point with the ORPOTrainer sample and locate the existing chunked_jsd_loss implementation. Review how teacher and student models are handled under FSDP, then define the required integration behavior and validation before proposing an implementation; the issue does not name specific files or tests.

Written by the indexing model from the issue text.