fix: fix single controller grad_norm

[yuki-97]

Bug

SC (Single Controller) split-API path had grad_norm at ~1/2 of the sync path from step 2 onward on the Megatron backend; reward regressed correspondingly.

Fix

Add torch.cuda.synchronize() between finish_grad_sync() and optimizer.step() in _finish_train_step_body, so the DP reduce dispatched on the comm stream completes before optimizer.step reads main_grad.

Test Results

• baseline (old code path)
  • yellow: didn't enable validate and checkpointing to keep the same test settings.
  • purple: with val_period=10, which will reset vLLM's prefix cache every 10 steps.
• blue: before this PR, grad_norm is ~1/2x smaller and reward is a bit lower.
• red: after this PR.

Note: Only the purple one will reset vLLM's prefix cache every 10 steps, other runs didn't reset for the whole run.

[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.

[yuki-97]

/ok to test c2feb15