Refactor chunked preference functions and distillation base class

[shivam15s]

Summary

Remove redundant code by refactoring

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[hebiao064]

I wonder is it full_target or actually target_chunk?

From the fused function, we are feeding into target_chunk

        def fused_fwd_bwd(
            input_chunk, target_chunk, ref_input_chunk, preference_labels_chunk
        ):
            """
            Fused forward and backward pass for a chunk of input and target.
            """
            if bias is not None:
                return torch.func.grad_and_value(
                    compute_loss, argnums=(0, 1, 3), has_aux=True
                )(
                    input_chunk,
                    weight,
                    target_chunk,
                    bias,
                    ref_input_chunk=ref_input_chunk,
                    preference_labels=preference_labels_chunk,
                )