[3/N] Achieve fault tolerance at the DP level (#11657)

Co-authored-by: UNIDY <unidy2002@outlook.com> Co-authored-by: Hank Han <hanhan7630@outlook.com>
2026-01-20 18:47:08 +08:00
parent d50dcd9b61
commit f7a5e425c3
8 changed files with 143 additions and 67 deletions
@@ -2238,6 +2238,27 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                reinit_attn_backend,
                split_forward_count,
            )
+            elastic_ep_state = ElasticEPStateManager.instance()
+            if (
+                elastic_ep_state is not None
+                and not elastic_ep_state.is_active_equal_last()
+            ):
+                elastic_ep_state.snapshot_active_to_last()
+                elastic_ep_state.sync_active_to_cpu()
+                logging.info("EPLB due to rank faults")
+                gen = self.eplb_manager.rebalance()
+                while True:
+                    try:
+                        next(gen)
+                    except StopIteration:
+                        break
+                output = self._forward_raw(
+                    forward_batch,
+                    skip_attn_backend_init,
+                    pp_proxy_tensors,
+                    reinit_attn_backend,
+                    split_forward_count,
+                )
        output.expert_distribution_metrics = recorder_outputs.get("metrics")

        # Copy cached routing experts' buffers back to CPU cache