Fix DeepCompile profile metadata backfill

deepspeed/compile/inductor.py

@@ -22,6 +22,23 @@
 from .partitioner import get_wrapped_partitioner
 
 
+def _get_graphsafe_run_with_rng_state():
+    try:
+        from torch._prims import rng_prims
+    except ImportError:
+        return None
+    return getattr(rng_prims, "graphsafe_run_with_rng_state", None)
+
+
+def _register_graphsafe_rng_state_no_reuse(register_fallback_no_reuse):
+    graphsafe_run_with_rng_state = _get_graphsafe_run_with_rng_state()
+    if graphsafe_run_with_rng_state is None:
+        return False
+
+    register_fallback_no_reuse(graphsafe_run_with_rng_state, never_reuse_output=True)
+    return True
+
+
 def patch_compiler(original_compiler, dc_compiler, z3_partition: bool, graph_id, graph_param_manager, bwd: bool):
 
     def wrapped_compiler(gm, fake_inputs):
@@ -243,6 +260,7 @@ def register_fallback_no_reuse(op_overload,
                                force_free_input=True)
     register_fallback_no_reuse(torch.ops.dc.free_tensors.default, never_reuse_input=True, never_reuse_output=True)
     register_fallback_no_reuse(torch.ops.dc.end_backward.default, never_reuse_input=True, never_reuse_output=False)
+    _register_graphsafe_rng_state_no_reuse(register_fallback_no_reuse)
 
     if not hasattr(Scheduler, "is_dc_patched") or not Scheduler.is_dc_patched:
         Scheduler.is_dc_patched = True

deepspeed/compile/passes/prefetch.py

@@ -12,6 +12,7 @@
 import deepspeed.comm as dist
 
 from ..profilers.comm_profile import create_predictor
+from ..profilers.graph_profile import is_profile_incomplete
 from ..graph_param import DSGraphParamManager
 
 NAME = "prefetch"
@@ -38,6 +39,11 @@ def schedule_prefetch(gm: GraphModule, graph_id: int, graph_order: List[Tuple[in
                       create_inputs_fn, mem_budget: float, param_manager: DSGraphParamManager,
                       bwd: bool) -> GraphModule:
 
+    profile_graph = profiling_results[graph_id].bwd_graph if bwd else profiling_results[graph_id].fwd_graph
+    if is_profile_incomplete(profile_graph):
+        print_rank_0(f"schedule_prefetch graph_id={graph_id} incomplete profiling data; skipping prefetch")
+        return gm
+
     max_mem = get_accelerator().total_memory() * (1 - MARGIN)
     vals_to_bcast = torch.tensor([max_mem], device=torch.device(get_accelerator().current_device()))
     dist.all_reduce(vals_to_bcast, dist.ReduceOp.MIN)

deepspeed/compile/passes/selective_gather.py

@@ -15,6 +15,7 @@
 
 from ..util import get_deepcompile_handle
 from ..graph_param import DSGraphParamManager
+from ..profilers.graph_profile import is_profile_incomplete
 
 NAME = "selective_gather"
 
@@ -27,6 +28,18 @@ def print_rank_0(message):
     log_dist(message, ranks=[0])
 
 
+def _maybe_update_size_from_profile(ds_id_to_size: Dict[int, int], ds_id: int, tensor_size: int) -> None:
+    if tensor_size > 0:
+        ds_id_to_size[ds_id] = tensor_size
+
+
+def _time_per_byte(ds_id_to_time: Dict[int, float], ds_id_to_size: Dict[int, int], ds_id: int) -> float:
+    size = ds_id_to_size.get(ds_id, 0)
+    if size <= 0:
+        return 0.0
+    return ds_id_to_time[ds_id] / size
+
+
 def _compute_persistence_budget(all_graph_mem_records: List[List[Tuple[str, int, int, int]]], total_mem: int,
                                 mem_margin: float) -> Dict[str, int]:
     usable_mem = int(total_mem * (1 - mem_margin))
@@ -54,6 +67,10 @@ def _compute_persistence_budget(all_graph_mem_records: List[List[Tuple[str, int,
     }
 
 
+def _profile_result_incomplete(prof) -> bool:
+    return is_profile_incomplete(prof.fwd_graph) or is_profile_incomplete(prof.bwd_graph)
+
+
 def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int, bool]], profiling_results,
                      create_inputs_fn, mem_budget: float, param_manager: DSGraphParamManager,
                      bwd: bool) -> GraphModule:
@@ -72,6 +89,14 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int
     if last_backward_graph_id is None or graph_id != last_backward_graph_id:
         return gm
 
+    incomplete_profile_ids = [
+        profile_graph_id for profile_graph_id, prof in profiling_results.items() if _profile_result_incomplete(prof)
+    ]
+    if incomplete_profile_ids:
+        print_rank_0(f"selective_gather incomplete profiling data for graph_ids={incomplete_profile_ids}; "
+                     "skipping persistence update")
+        return gm
+
     all_graph_mem_records = []
     for profile_graph_id, prof in profiling_results.items():
         all_graph_mem_records.extend([prof.fwd_mem, prof.bwd_mem])
@@ -106,7 +131,7 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int
         for n in profile.fwd_graph.nodes:
             if n.target == torch.ops.dc.allgather_param.default:
                 assert "tensor_size" in n.meta
-                ds_id_to_size[n.args[2]] = n.meta["tensor_size"]
+                _maybe_update_size_from_profile(ds_id_to_size, n.args[2], n.meta["tensor_size"])
                 assert "device_time" in n.meta
                 ds_id_to_time[n.args[2]] += n.meta["device_time"]
 
@@ -117,12 +142,12 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int
             for n in profile.bwd_graph.nodes:
                 if n.target == torch.ops.dc.allgather_param.default:
                     assert "tensor_size" in n.meta
-                    ds_id_to_size[n.args[2]] = n.meta["tensor_size"]
+                    _maybe_update_size_from_profile(ds_id_to_size, n.args[2], n.meta["tensor_size"])
                     assert "device_time" in n.meta
                     ds_id_to_time[n.args[2]] += n.meta["device_time"]
 
     ds_ids = [ds_id for ds_id in ds_id_to_size if ds_id not in persistent_ds_ids]
-    ds_ids.sort(key=lambda ds_id: ds_id_to_time[ds_id] / ds_id_to_size[ds_id], reverse=True)
+    ds_ids.sort(key=lambda ds_id: _time_per_byte(ds_id_to_time, ds_id_to_size, ds_id), reverse=True)
 
     # print(f"ds_id_to_size={ds_id_to_size}")
     # print(f"ds_id_to_time={ds_id_to_time}")

deepspeed/compile/profilers/graph_profile.py

@@ -8,7 +8,7 @@
 import statistics
 
 import torch
-from torch.fx import GraphModule, Interpreter
+from torch.fx import Graph, GraphModule, Interpreter
 from torch.fx.node import map_aggregate
 
 try:
@@ -54,6 +54,43 @@ def _node_size(out):
     return sum([v.element_size() * v.numel() for v in tree_leaves(out) if torch.is_tensor(v)])
 
 
+_PROFILE_META_DEFAULTS = {
+    "device_time": 0.0,
+    "wall_time": 0.0,
+    "tensor_size": 0,
+    "alloc_mem": 0,
+    "max_mem": 0,
+}
+_PROFILE_INCOMPLETE_ATTR = "_deepcompile_profile_incomplete"
+_PROFILE_INCOMPLETE_META_KEY = "deepcompile_profile_incomplete"
+
+
+def _mark_profile_incomplete(graph: Graph):
+    setattr(graph, _PROFILE_INCOMPLETE_ATTR, True)
+    for node in graph.nodes:
+        node.meta[_PROFILE_INCOMPLETE_META_KEY] = True
+
+
+def is_profile_incomplete(graph: Graph):
+    if graph is None:
+        return False
+    if getattr(graph, _PROFILE_INCOMPLETE_ATTR, False):
+        return True
+    return any(node.meta.get(_PROFILE_INCOMPLETE_META_KEY, False) for node in graph.nodes)
+
+
+def _has_missing_profile_metadata(graph: Graph):
+    return any(key not in node.meta for node in graph.nodes for key in _PROFILE_META_DEFAULTS)
+
+
+def _backfill_missing_profile_metadata(graph: Graph, profile_complete: bool = True):
+    if not profile_complete or _has_missing_profile_metadata(graph):
+        _mark_profile_incomplete(graph)
+    for node in graph.nodes:
+        for key, default in _PROFILE_META_DEFAULTS.items():
+            node.meta.setdefault(key, default)
+
+
 def _get_mem_usage_out_of_torch():
 
     adjust = 0
@@ -100,6 +137,7 @@ def run(self, *args) -> Any:
         returns: The output of the graph. Tensor in the output is real tensors.
         """
         return_val = None
+        profile_complete = True
         try:
             assert _all_real_if_tensor(args), "Inputs must be real tensors"
             self.nz3.enable_profiling(True)
@@ -109,11 +147,18 @@ def run(self, *args) -> Any:
                     self.mem_usage_out_of_torch = _get_mem_usage_out_of_torch()
                     return_val = super().run(*args)
         except Exception as e:
+            profile_complete = False
             msg = e.msg if "msg" in dir(e) else str(e)
-            print(f"Profiling error {msg}")
+            if not self.distributed or dist.get_rank() == 0:
+                print(f"DeepCompile profiling failed; using default profile metadata for incomplete nodes: {msg}")
         finally:
-            self.nz3.clear_all_gathered_params()
-            self.nz3.enable_profiling(False)
+            try:
+                self.nz3.clear_all_gathered_params()
+            finally:
+                try:
+                    self.nz3.enable_profiling(False)
+                finally:
+                    _backfill_missing_profile_metadata(self.graph, profile_complete=profile_complete)
         return return_val
 
     def run_node(self, n: torch.fx.Node) -> Any:

tests/unit/compile/test_list_schedule.py

@@ -4,13 +4,19 @@
 # DeepSpeed Team
 
 import operator
+from types import SimpleNamespace
 
 import pytest
 import torch
-from torch.fx import Graph
+from torch.fx import Graph, GraphModule
 
 import deepspeed.compile.util as compile_util
+from deepspeed.compile import inductor as inductor_mod
 from deepspeed.compile import list_schedule as schedule_mod
+from deepspeed.compile.passes import prefetch as prefetch_mod
+from deepspeed.compile.passes import selective_gather as selective_gather_mod
+from deepspeed.compile.profilers import ProfilingResult
+from deepspeed.compile.profilers.graph_profile import _backfill_missing_profile_metadata, is_profile_incomplete
 
 _DC_LIBRARIES = []
 
@@ -31,6 +37,8 @@ def _define_dc_ops():
             "wait_allgather(Tensor(a) a, int graph_id, int id) -> Tensor(a)",
             "release_param(Tensor(a) a, int graph_id, int id, int n_users) -> Tensor(a)",
             "reduce_grad(Tensor a, int graph_id, int id) -> Tensor",
+            "free_tensors(Tensor[] tensors) -> ()",
+            "end_backward(Tensor[] tensors, int graph_id, bool release_reduce_buckets = True) -> ()",
     ):
         try:
             lib.define(schema)
@@ -49,7 +57,8 @@ def stub_deepcompile_ops(monkeypatch):
 
 def _with_meta(node, tensor_size=0, device_time=0):
     node.meta["tensor_size"] = tensor_size
-    node.meta["device_time"] = device_time
+    if device_time is not None:
+        node.meta["device_time"] = device_time
     return node
 
 
@@ -199,3 +208,137 @@ def test_fast_free_schedule_keeps_single_allgather_release_order():
     assert names.index(ag.name) < names.index(wait.name)
     assert names.index(wait.name) < names.index(use.name)
     assert names.index(use.name) < names.index(release.name)
+
+
+def test_profile_backfill_makes_partial_profile_safe_for_profile_dependent_passes(monkeypatch):
+    graph = Graph()
+
+    param = _placeholder(graph, "partial_profile_param")
+    ag = _allgather(graph, param, 90, "partial_profile", device_time=None)
+    wait = _wait(graph, ag, 90, "partial_profile")
+    use = _neg(graph, wait, "partial_profile_use", device_time=None)
+    release = _release(graph, use, 90, "partial_profile")
+
+    ag.meta.pop("tensor_size", None)
+    for node in (ag, use):
+        node.meta.pop("wall_time", None)
+        node.meta.pop("alloc_mem", None)
+        node.meta.pop("max_mem", None)
+
+    graph.output((release, ))
+    graph.lint()
+
+    _backfill_missing_profile_metadata(graph)
+    assert is_profile_incomplete(graph)
+
+    for node in graph.nodes:
+        if node in (ag, use):
+            assert node.meta["device_time"] == 0.0
+        else:
+            assert "device_time" in node.meta
+        assert "wall_time" in node.meta
+        assert "tensor_size" in node.meta
+        assert "alloc_mem" in node.meta
+        assert "max_mem" in node.meta
+    assert ag.meta["tensor_size"] == 0
+
+    names = _scheduled_names(graph)
+    assert names.index(ag.name) < names.index(wait.name)
+    assert names.index(wait.name) < names.index(use.name)
+    assert names.index(use.name) < names.index(release.name)
+
+    class FakeAccelerator:
+
+        def current_device(self):
+            return "cpu"
+
+        def total_memory(self):
+            return 1024
+
+        def available_memory(self):
+            return 1024
+
+    fake_ds_param = SimpleNamespace(numel=7,
+                                    dtype=torch.float16,
+                                    param=SimpleNamespace(ds_persist=False, ds_shape=(1, )))
+    fake_param_manager = {
+        0: SimpleNamespace(params={"partial_profile_param": fake_ds_param}, ds_ids={"partial_profile_param": 90})
+    }
+    profiling_results = {
+        0: ProfilingResult(fwd_graph=graph, bwd_graph=None, fwd_mem=[("profiled_before_abort", 0, 0, 0)])
+    }
+    gm = GraphModule(torch.nn.Module(), graph)
+    logs = []
+    prefetch_logs = []
+    persisted = []
+
+    monkeypatch.setattr(prefetch_mod, "print_rank_0", lambda message: prefetch_logs.append(message))
+    assert prefetch_mod.schedule_prefetch(gm,
+                                          graph_id=0,
+                                          graph_order=[(0, True)],
+                                          profiling_results=profiling_results,
+                                          create_inputs_fn=lambda: (),
+                                          mem_budget=0,
+                                          param_manager=fake_param_manager,
+                                          bwd=False) is gm
+    assert any("incomplete profiling data" in message for message in prefetch_logs)
+
+    monkeypatch.setattr(selective_gather_mod, "print_rank_0", lambda message: logs.append(message))
+    monkeypatch.setattr(selective_gather_mod, "get_accelerator", lambda: FakeAccelerator())
+    monkeypatch.setattr(selective_gather_mod, "get_deepcompile_handle",
+                        lambda: SimpleNamespace(set_persistent=persisted.append))
+    monkeypatch.setattr(selective_gather_mod.dist, "all_reduce", lambda *args, **kwargs: None)
+
+    selective_gather_mod.selective_gather(gm,
+                                          graph_id=0,
+                                          graph_order=[(0, True)],
+                                          profiling_results=profiling_results,
+                                          create_inputs_fn=lambda: (),
+                                          mem_budget=0,
+                                          param_manager=fake_param_manager,
+                                          bwd=True)
+    assert persisted == []
+    assert any("incomplete profiling data" in message for message in logs)
+
+
+def test_graphsafe_rng_state_outputs_are_registered_no_reuse():
+    graphsafe_run_with_rng_state = inductor_mod._get_graphsafe_run_with_rng_state()
+    if graphsafe_run_with_rng_state is None:
+        pytest.skip("graphsafe_run_with_rng_state is unavailable in this torch build")
+
+    calls = []
+
+    def fake_register(op_overload, **kwargs):
+        calls.append((op_overload, kwargs))
+
+    assert inductor_mod._register_graphsafe_rng_state_no_reuse(fake_register)
+    assert calls == [(graphsafe_run_with_rng_state, {"never_reuse_output": True})]
+
+
+def test_register_custom_ops_includes_graphsafe_rng_state_no_reuse(monkeypatch):
+    graphsafe_run_with_rng_state = inductor_mod._get_graphsafe_run_with_rng_state()
+    if graphsafe_run_with_rng_state is None:
+        pytest.skip("graphsafe_run_with_rng_state is unavailable in this torch build")
+
+    _define_dc_ops()
+    registered_ops = []
+
+    def fake_add_needs_realized_inputs(_op_overload):
+        return None
+
+    def fake_register_lowering(op_overload, **_kwargs):
+
+        def record_handler(handler):
+            registered_ops.append(op_overload)
+            return handler
+
+        return record_handler
+
+    monkeypatch.setattr(inductor_mod, "add_needs_realized_inputs", fake_add_needs_realized_inputs)
+    monkeypatch.setattr(inductor_mod, "register_lowering", fake_register_lowering)
+    monkeypatch.setattr(inductor_mod, "fallbacks", set())
+    monkeypatch.setattr(inductor_mod.Scheduler, "is_dc_patched", True, raising=False)
+
+    inductor_mod.register_custom_ops()
+
+    assert graphsafe_run_with_rng_state in registered_ops