fix: Fix DAPO crash due to discard generation weights

[guyueh1]

What does this PR do ?

Fix an issue introduced by #2495 : when using dynamic sampling, at end of generation phase, we cannot simply discard generation weights and mark it stale, because if the batch is decided to be incomplete for training, we directly jump to next step and start generation again, in that case marking it stale will trigger a refit which will crash.

This is just a temporal patch, but it does not enable the intended optimization of #2495 for DAPO; I will propose a proper fix in another PR.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[guyueh1]

/ok to test 824462d

[guyueh1]

/ok to test 232d44f

[sharonyu-115]

Hi @guyueh1, I ran into the problem to be fixed by this PR when testing #2224. Then I found #2587 and #2614 seem to address the same problem, and the latter two are quite similar. Looks like the fixes need to be consolidated?

@jQizhang fyi.

[guyueh1]

Hi @guyueh1, I ran into the problem to be fixed by this PR when testing #2224. Then I found #2587 and #2614 seem to address the same problem, and the latter two are quite similar. Looks like the fixes need to be consolidated?

@jQizhang fyi.

@sharonyu-115
they fix the same problem but in different ways, using this PR, dapo will basically behave the same as before the breaking commit; using 2614 or 2587, it fixes the breaking commit but will slightly regress the perf (more invocation of refit) but saves cpu memory. The best fix is to not sleep the generation engine at all if the dynamic sampling has not finished a complete batch, but temporally both this or that can work. I'ok with whichever is merged.