[CP]: Support CP V2 Strategy for dsv4 (#33532)

2026-08-07 14:03:27 -07:00
parent 8e7d361def
commit 1480687cff
10 changed files with 161 additions and 59 deletions
@@ -374,12 +374,18 @@ class EagerRunner(BaseRunner):
        hidden_states = cp_gather_after_forward(
            hidden_states, forward_batch, torch.cuda.current_stream()
        )
+        logits_kwargs = {}
+        # DSV4 returns (hidden_states, hidden_states_before_norm) from its model body.
+        if isinstance(hidden_states, tuple):
+            hidden_states, hidden_states_before_norm = hidden_states
+            logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
        return model.logits_processor(
            forward_batch.input_ids,
            hidden_states,
            model.lm_head,
            forward_batch,
            aux_hidden_states,
+            **logits_kwargs,
        )

    def _execute_idle(