What does this PR do?

This PR implements fused losses for alignment. #710
It reduces the memory required for loss calculation to a small constant amount.

ChangeLog:

• added the option use_fused_kernels
• monkey patch to make model.forward return last_hidden_state and not calculate logits
• Added FusedLinearForPPO to verl/utils/experimental/torch_functional.py

Usage

Simply add the following option

actor_rollout_ref.model.use_fused_kernels=True

Before submitting

• Did you read the Contribute Guide and finish the code format check?
• Did you make sure to update the documentations with your changes in the docs especially for breaking config etc?
• Did you write any test cases if neccessary? Please add CI tests to your new feature.

Additional Info:

• The current implementation uses chunking to reduce the memory consumption to a constant value.
  • It works by splitting the loss calculations into chunks of 512 tokens. Calculating the log_probs / entropy values / gradients for each chunk and accumulating them.
  • However the current implementation can be slow. It processes each chunk sequentially in a python for loop.
  • In the future we should consider converting the fused functions into triton or some other JIT solution.
• Compared to FusedPPOLossFunction, optimizing hidden_states -> entropy & log_probs is much better for algorithm developers as the memory heavy part is optimized away for them and they are free to combine the values for their own custom loss functions.