diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index 1c2624b..04fe271 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -31,7 +31,10 @@ transfer stream) and targets only decrease. Compute is claimed a whole node at a time, so the gather that stops in front of a long kernel leaves most of it unused; that remainder carries to the next gather instead of being discarded, which is -what keeps a 95us collective from spending a 33ms attention. All moves are +what keeps a 95us collective from spending a 33ms attention. The remainder is +dropped once the sweep reaches a gather that already sits upstream of the node +holding it -- that compute is outside such a gather's window, and counting it +would report the gather as covered and leave it unmoved. All moves are applied in one stable-sort rebuild and validated once (``_validate_full``) -- the Inductor driver does NOT repair the returned order, so it must be a valid topological order. @@ -323,18 +326,25 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: targets: dict = {} # launch -> target index (in original order space) compute_idx = len(order) # scan compute strictly below this carry = 0.0 # runtime the previous gather left unspent in its boundary node + # Index of that boundary node: its leftover is only real for gathers that + # still sit at or after it. + carry_idx = len(order) for launch, group, fc_idx, comm_runtime, lower in reversed(plans): cur = index_of[launch] # Start just before the launch, but no later than where the previous # (later) gather already consumed compute down to. compute_idx = min(compute_idx, cur) need = comm_runtime * self.comm_overlap_window_scale + self.comm_overlap_window_margin_ns + if cur < carry_idx: + carry = 0.0 + carry_idx = compute_idx acc = carry_in = carry t = compute_idx while acc < need and t > lower: s = order[t - 1] if self._is_compute(s): acc += self._cost(s) + carry_idx = t - 1 # last node claimed; holds this gather's leftover t -= 1 # target == cur means no upstream compute left (graph head or previous # gather claimed it); target >= lower keeps real producers before it. @@ -351,14 +361,16 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: # comm = the gather's runtime it needs to hide # carry_in = capacity inherited from the later gather's boundary node # (target==cur with a large carry_in means it was already - # covered and did not have to move at all) + # covered and did not have to move at all); 0 when this + # gather already sat upstream of that boundary node + # carry_idx = boundary node this gather leaves its own remainder in # acc_upstream = carry_in plus the compute found in [target, cur] # verdict = hidden (acc>=need) | COMPUTE-LIMITED (ran out of upstream # compute before covering comm -- i.e. hit `lower` or the # previous gather's placement first) magi_logger.debug( "FSDP overlap placement: launch %s(%s) cur=%d -> target=%d fc=%d lower=%d | " - "comm=%.1fus carry_in=%.1fus acc_upstream=%.1fus need=%.1fus %s", + "comm=%.1fus carry_in=%.1fus carry_idx=%d acc_upstream=%.1fus need=%.1fus %s", launch.get_name(), getattr(_leaf_collective_node(launch), "op_overload", "?"), cur, @@ -367,6 +379,7 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: lower, comm_runtime / 1e3, carry_in / 1e3, + carry_idx, acc / 1e3, need / 1e3, "hidden" if acc >= need else "COMPUTE-LIMITED",