Support true on-policy (#12058)

2025-10-25 10:23:42 +08:00
parent 649949807f
commit 20bd2271e2
16 changed files with 151 additions and 11 deletions
--- a/python/sglang/srt/layers/logits_processor.py
+++ b/python/sglang/srt/layers/logits_processor.py
@@ -593,6 +593,11 @@ class LogitsProcessor(nn.Module):
                    None,  # bias
                    True,  # is_vnni
                )
+            elif get_global_server_args().rl_on_policy_target == "fsdp":
+                # Due to tie-weight, we may not be able to change lm_head's weight dtype
+                logits = torch.matmul(
+                    hidden_states.bfloat16(), lm_head.weight.T.bfloat16()
+                )
            else:
                logits = torch.matmul(
                    hidden_states.to(lm_head.weight.dtype), lm_head.weight.T