[Refactor] fused kernel in forward

[mingruimingrui]

Checklist Before Starting

• Search for similar PR(s).

What does this PR do?

Shifts fused_linear_for_ppo into model.forward for FSDP

High-Level Design

Self explaining

Specific Changes

• Update monkey patch to return log_probs and entropy instead of last_hidden_state.

API

No changes

Usage Example

actor_rollout_ref.model.use_fused_kernels=True

Test

Additional Info.

• This is to fix Running error on PPO training on GSM8K dataset example. "RuntimeError: setStorage:" #1565
• The original bug arises because we tried to access model.lm_head.weight from outside of the FSDP wrapped context.

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks.
• Add [BREAKING] to the PR title if it breaks any API.
• Update the documentation about your changes in the docs.
• Add CI test(s) if necessary.

[mingruimingrui]

Hi @ETOgaosion, I've made a fix for the fused kernel bug that we experienced when using 1 GPU. It's strange that we didn't encounter this when using more than 1.

[ETOgaosion]

Seems that currently there is no CI enables fused_kernel, could you enable some tests and disable in other tests to test both cases?

[mingruimingrui]

@ETOgaosion My bad, just added 😂

[ETOgaosion]

@mingruimingrui #1634 fixed the CI, could you rebase main and retry please?

[mingruimingrui]

Done, trouble you to trigger CI again