diff --git a/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py b/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py index c318e36..b93029c 100644 --- a/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py +++ b/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py @@ -25,6 +25,9 @@ class P0Worker(Worker): original = model_runner._determine_batch_execution_and_padding def traced_determine(*args: Any, **kwargs: Any): # type: ignore[no-untyped-def] + local_unpadded_rows = kwargs.get("num_tokens") + if local_unpadded_rows is None and len(args) > 1: + local_unpadded_rows = args[1] output = original(*args, **kwargs) self._p0_batch_epoch += 1 scheduler_output = kwargs.get("scheduler_output") @@ -38,6 +41,7 @@ class P0Worker(Worker): **self._p0_rank_fields(), **plan_summary(scheduler_output), cudagraph_mode=getattr(cudagraph_mode, "value", str(cudagraph_mode)), + local_unpadded_rows=local_unpadded_rows, physical_batch_rows=getattr(batch_desc, "num_tokens", None), should_ubatch=should_ubatch, rows_across_dp=rows_across_dp, @@ -54,7 +58,9 @@ class P0Worker(Worker): "global_rank": getattr(self, "rank", None), "local_rank": getattr(self, "local_rank", None), "data_parallel_rank": getattr(config, "data_parallel_rank", None), - "tensor_parallel_rank": getattr(config, "tensor_parallel_rank", None), + # Worker.rank is the tensor-parallel rank within a DP replica in + # this vLLM executor. ParallelConfig does not expose it here. + "tensor_parallel_rank": getattr(self, "rank", None), "expert_parallel_size": getattr(config, "expert_parallel_size", None), "tensor_parallel_size": getattr(config, "tensor_parallel_size", None), "data_parallel_size": getattr(config, "data_parallel_size", None), diff --git a/scripts/collectivespec/summarize_p0_phase_trace.py b/scripts/collectivespec/summarize_p0_phase_trace.py index c21aca6..ba9349d 100644 --- a/scripts/collectivespec/summarize_p0_phase_trace.py +++ b/scripts/collectivespec/summarize_p0_phase_trace.py @@ -105,48 +105,264 @@ def rank_key(event: dict[str, Any]) -> str | None: return f"dp{dp_rank}/local{local_rank}" -def sequence_agreement(worker_events: list[dict[str, Any]]) -> dict[str, Any]: - begins = [event for event in worker_events if event.get("event") == "batch_execution_plan"] +def _event_order(event: dict[str, Any], epoch_field: str) -> tuple[int, float]: + epoch = integer(event.get(epoch_field)) + monotonic_s = event.get("monotonic_s") + return ( + epoch if epoch is not None else -1, + float(monotonic_s) if isinstance(monotonic_s, (int, float)) else 0.0, + ) + + +def _canonical(value: Any) -> str: + return json.dumps(value, ensure_ascii=True, separators=(",", ":"), sort_keys=True) + + +def _target_event(event: dict[str, Any]) -> bool: + """Exclude profile/dummy calls with an empty SchedulerOutput. + + The worker hook also wraps vLLM's profile and DP dummy runs. Those calls + have a physical batch descriptor but no target request, so their ranks and + counts are not a valid collective-plan invariant. + """ + + return ( + (integer(event.get("request_count")) or 0) > 0 + and (integer(event.get("total_scheduled_rows")) or 0) > 0 + ) + + +def _target_signature(event: dict[str, Any]) -> str: + return _canonical( + { + "cudagraph_mode": event.get("cudagraph_mode"), + "ordered_plan_digest": event.get("ordered_plan_digest"), + "physical_batch_rows": event.get("physical_batch_rows"), + "rows_across_dp": event.get("rows_across_dp"), + "scheduled_dp_metadata": event.get("scheduled_dp_metadata"), + "semantic_plan_digest": event.get("semantic_plan_digest"), + "should_ubatch": event.get("should_ubatch"), + "total_scheduled_rows": event.get("total_scheduled_rows"), + } + ) + + +def _target_metadata_check(event: dict[str, Any]) -> tuple[bool, int | None, bool | None]: + """Check the local, post-DP-coordination shape invariants.""" + + metadata = event.get("scheduled_dp_metadata") + dp_rank = integer(event.get("data_parallel_rank")) + dp_size = integer(event.get("data_parallel_size")) + total_rows = integer(event.get("total_scheduled_rows")) + physical_rows = integer(event.get("physical_batch_rows")) + rows_across_dp = event.get("rows_across_dp") + if ( + not isinstance(metadata, dict) + or dp_rank is None + or dp_size is None + or dp_rank < 0 + or dp_rank >= dp_size + or total_rows is None + or physical_rows is None + or not isinstance(rows_across_dp, list) + ): + return False, None, None + raw_rows = metadata.get("num_tokens_per_rank") + padded_rows = metadata.get("num_tokens_across_dp") + if ( + not isinstance(raw_rows, list) + or not isinstance(padded_rows, list) + or len(raw_rows) != dp_size + or len(padded_rows) != dp_size + or len(rows_across_dp) != dp_size + or not all(integer(value) is not None and integer(value) >= 0 for value in raw_rows) + or not all(integer(value) is not None and integer(value) >= 0 for value in padded_rows) + or not all(integer(value) is not None and integer(value) >= 0 for value in rows_across_dp) + ): + return False, None, None + raw = [int(value) for value in raw_rows] + padded = [int(value) for value in padded_rows] + coordinated = [int(value) for value in rows_across_dp] + valid = ( + raw[dp_rank] == total_rows + and physical_rows == coordinated[dp_rank] + and padded == coordinated + and all(padded_value >= raw_value for padded_value, raw_value in zip(padded, raw)) + ) + return valid, padded[dp_rank] - raw[dp_rank], len(set(raw)) > 1 + + +def _dp_pair_key(event: dict[str, Any]) -> tuple[str, str] | None: + """Return the TP-shard key shared by the two actual DP peers. + + New P0 logs carry ``tensor_parallel_rank``. The v2 artifact predates that + field, where ``Worker.rank`` was logged as ``global_rank`` and is the same + TP-shard identifier within each DP replica. The fallback is explicit in + the output so it cannot be confused with a true global rank. + """ + + tp_rank = integer(event.get("tensor_parallel_rank")) + if tp_rank is not None: + return f"tp{tp_rank}", "tensor_parallel_rank" + legacy_worker_rank = integer(event.get("global_rank")) + if legacy_worker_rank is not None: + return f"legacy-worker-rank{legacy_worker_rank}", "legacy_worker_rank" + local_rank = integer(event.get("local_rank")) + tp_size = integer(event.get("tensor_parallel_size")) + if local_rank is not None and tp_size is not None and tp_size > 0: + return f"local-rank-mod-tp{local_rank % tp_size}", "local_rank_mod_tp_size" + return None + + +def _dp_coordination_signature(event: dict[str, Any]) -> str: + metadata = event.get("scheduled_dp_metadata") + if not isinstance(metadata, dict): + metadata = {} + return _canonical( + { + "all_decode": metadata.get("all_decode"), + "all_prefill": metadata.get("all_prefill"), + "cudagraph_mode": metadata.get("cudagraph_mode"), + "input_fits_in_drafter": metadata.get("input_fits_in_drafter"), + "is_prompt_batch": metadata.get("is_prompt_batch"), + "num_reqs_per_rank": metadata.get("num_reqs_per_rank"), + "num_tokens_across_dp": metadata.get("num_tokens_across_dp"), + "num_tokens_per_rank": metadata.get("num_tokens_per_rank"), + "rows_across_dp": event.get("rows_across_dp"), + "should_ubatch": metadata.get("should_ubatch"), + } + ) + + +def dp_pair_coordination(target_events: list[dict[str, Any]]) -> dict[str, Any]: + """Compare scalar DP coordination seen by the two peers of each TP shard.""" + + grouped: dict[str, dict[int, dict[int, dict[str, Any]]]] = defaultdict( + lambda: defaultdict(dict) + ) + sources: dict[str, set[str]] = defaultdict(set) + expected_sizes: dict[str, set[int]] = defaultdict(set) + for event in target_events: + key = _dp_pair_key(event) + dp_rank = integer(event.get("data_parallel_rank")) + epoch = integer(event.get("batch_phase_epoch")) + dp_size = integer(event.get("data_parallel_size")) + if key is None or dp_rank is None or epoch is None: + continue + group, source = key + grouped[group][dp_rank][epoch] = event + sources[group].add(source) + if dp_size is not None: + expected_sizes[group].add(dp_size) + + groups: dict[str, Any] = {} + for group, by_dp_rank in grouped.items(): + dp_ranks = sorted(by_dp_rank) + epoch_sets = [set(by_epoch) for by_epoch in by_dp_rank.values()] + shared_epochs = sorted(set.intersection(*epoch_sets)) if epoch_sets else [] + signatures = { + dp_rank: { + epoch: _dp_coordination_signature(event) + for epoch, event in by_epoch.items() + } + for dp_rank, by_epoch in by_dp_rank.items() + } + shared_equal = all( + len({signatures[dp_rank][epoch] for dp_rank in dp_ranks}) == 1 + for epoch in shared_epochs + ) + expected_values = expected_sizes.get(group, set()) + expected_dp_size = next(iter(expected_values)) if len(expected_values) == 1 else None + groups[group] = { + "grouping_source": sorted(sources[group]), + "data_parallel_ranks": dp_ranks, + "expected_data_parallel_size": expected_dp_size, + "target_epoch_count_by_dp_rank": { + str(dp_rank): len(by_dp_rank[dp_rank]) for dp_rank in dp_ranks + }, + "shared_target_epoch_count": len(shared_epochs), + "all_shared_epoch_signatures_identical": bool(shared_epochs) and shared_equal, + "all_expected_dp_peers_observed": expected_dp_size is not None + and len(dp_ranks) == expected_dp_size, + } + return { + "group_count": len(groups), + "groups": groups, + "all_expected_dp_peers_observed": bool(groups) + and all(value["all_expected_dp_peers_observed"] for value in groups.values()), + "all_shared_epoch_signatures_identical": bool(groups) + and all(value["all_shared_epoch_signatures_identical"] for value in groups.values()), + } + + +def target_execution_agreement(worker_events: list[dict[str, Any]]) -> dict[str, Any]: + batch_events = [event for event in worker_events if event.get("event") == "batch_execution_plan"] + target_events = [event for event in batch_events if _target_event(event)] per_rank: dict[str, list[dict[str, Any]]] = defaultdict(list) - per_dp: dict[str, set[str]] = defaultdict(set) - for event in begins: + per_replica: dict[str, set[str]] = defaultdict(set) + for event in target_events: rank = rank_key(event) - dp_rank = event.get("data_parallel_rank") - if rank is None: + dp_rank = integer(event.get("data_parallel_rank")) + if rank is None or dp_rank is None: continue per_rank[rank].append(event) - if dp_rank is not None: - per_dp[str(dp_rank)].add(rank) + per_replica[str(dp_rank)].add(rank) + sequences: dict[str, list[str]] = {} for rank, values in per_rank.items(): - values.sort(key=lambda item: (integer(item.get("worker_phase_epoch")) or -1, float(item.get("monotonic_s", 0.0)))) - sequences[rank] = [str(item.get("ordered_plan_digest")) for item in values] + values.sort(key=lambda event: _event_order(event, "batch_phase_epoch")) + sequences[rank] = [_target_signature(event) for event in values] counts = {rank: len(values) for rank, values in sequences.items()} - within_dp: dict[str, Any] = {} - for dp_rank, ranks in per_dp.items(): + + replicas: dict[str, Any] = {} + for dp_rank, ranks in per_replica.items(): rank_list = sorted(ranks) distinct = {tuple(sequences[rank]) for rank in rank_list} - within_dp[dp_rank] = { + phase_counts = [counts[rank] for rank in rank_list] + replicas[dp_rank] = { "ranks": rank_list, "phase_count_by_rank": {rank: counts[rank] for rank in rank_list}, + "phase_count_equal": bool(phase_counts) and len(set(phase_counts)) == 1, "sequence_distinct_count": len(distinct), - "identical_execution_sequence": len(distinct) == 1, + "identical_target_execution_sequence": len(distinct) == 1, } - count_values = list(counts.values()) + + metadata_valid_count = 0 + metadata_padding: list[int] = [] + raw_unequal_count = 0 + for event in target_events: + valid, padding, raw_unequal = _target_metadata_check(event) + metadata_valid_count += valid + if padding is not None: + metadata_padding.append(padding) + raw_unequal_count += bool(raw_unequal) return { - "worker_begin_event_count": len(begins), + "batch_execution_plan_count": len(batch_events), + "dummy_or_profile_event_count": len(batch_events) - len(target_events), + "target_event_count": len(target_events), "worker_rank_count": len(sequences), - "phase_count_by_rank": counts, - "phase_count": { - "n": len(count_values), - "min": min(count_values) if count_values else None, - "max": max(count_values) if count_values else None, - "distinct_value_count": len(set(count_values)), - "all_equal": bool(count_values) and len(set(count_values)) == 1, + "target_phase_count_by_rank": counts, + "within_dp_replica_tp_group": replicas, + "all_target_phase_counts_equal_within_dp_replica": bool(replicas) + and all(value["phase_count_equal"] for value in replicas.values()), + "all_target_sequences_identical_within_dp_replica": bool(replicas) + and all(value["identical_target_execution_sequence"] for value in replicas.values()), + "dp_metadata": { + "target_record_count": len(target_events), + "valid_record_count": metadata_valid_count, + "all_records_valid": len(target_events) > 0 + and metadata_valid_count == len(target_events), + "raw_dp_counts_unequal_record_count": raw_unequal_count, + "post_coordinate_minus_local_raw_rows": { + "n": len(metadata_padding), + "min": min(metadata_padding) if metadata_padding else None, + "max": max(metadata_padding) if metadata_padding else None, + "sum": sum(metadata_padding), + "distinct_value_count": len(set(metadata_padding)), + "non_negative": all(value >= 0 for value in metadata_padding), + }, }, - "within_data_parallel_replica": within_dp, - "all_within_dp_sequences_identical": bool(within_dp) - and all(value["identical_execution_sequence"] for value in within_dp.values()), + "dp_pair_coordination": dp_pair_coordination(target_events), } @@ -160,10 +376,11 @@ def summarize_cell(cell: Path, expected_requests: int, expected_tokens: int) -> event for event in schedule_events if integer(event.get("spec_request_count")) not in (None, 0) ] batch_events = [event for event in worker_events if event.get("event") == "batch_execution_plan"] - rows_across_dp = [event.get("rows_across_dp") for event in batch_events] + target_events = [event for event in batch_events if _target_event(event)] + rows_across_dp = [event.get("rows_across_dp") for event in target_events] physical_rows = [ integer(event.get("physical_batch_rows")) - for event in batch_events + for event in target_events if integer(event.get("physical_batch_rows")) is not None ] return { @@ -186,10 +403,9 @@ def summarize_cell(cell: Path, expected_requests: int, expected_tokens: int) -> ), }, "worker": { - **sequence_agreement(worker_events), - "batch_execution_plan_count": len(batch_events), - "rows_across_dp_present_count": sum(value is not None for value in rows_across_dp), - "physical_batch_rows": { + **target_execution_agreement(worker_events), + "target_rows_across_dp_present_count": sum(value is not None for value in rows_across_dp), + "target_physical_batch_rows": { "n": len(physical_rows), "min": min(physical_rows) if physical_rows else None, "max": max(physical_rows) if physical_rows else None, @@ -227,18 +443,46 @@ def data_sanity(cells: dict[str, dict[str, Any]], expected_gpu_count: int | None bool(cell.get("probe_integrity", {}).get("valid")) for cell in cells.values() ), "no_event_parse_errors": all(not cell.get("event_parse_errors") for cell in cells.values()), - "all_rank_phase_counts_equal": all( - bool(cell.get("worker", {}).get("phase_count", {}).get("all_equal")) + "all_target_phase_counts_equal_within_dp_replica": all( + bool( + cell.get("worker", {}).get( + "all_target_phase_counts_equal_within_dp_replica" + ) + ) for cell in cells.values() ), - "all_within_dp_sequences_identical": all( - bool(cell.get("worker", {}).get("all_within_dp_sequences_identical")) + "all_target_sequences_identical_within_dp_replica": all( + bool( + cell.get("worker", {}).get( + "all_target_sequences_identical_within_dp_replica" + ) + ) for cell in cells.values() ), - "all_dp_coordination_records_observed": all( - integer(cell.get("worker", {}).get("rows_across_dp_present_count")) - == integer(cell.get("worker", {}).get("batch_execution_plan_count")) - and integer(cell.get("worker", {}).get("batch_execution_plan_count", 0)) > 0 + "all_target_dp_metadata_valid": all( + bool(cell.get("worker", {}).get("dp_metadata", {}).get("all_records_valid")) + for cell in cells.values() + ), + "all_target_dp_coordination_records_observed": all( + integer(cell.get("worker", {}).get("target_rows_across_dp_present_count")) + == integer(cell.get("worker", {}).get("target_event_count")) + and integer(cell.get("worker", {}).get("target_event_count", 0)) > 0 + for cell in cells.values() + ), + "all_expected_dp_peers_observed": all( + bool( + cell.get("worker", {}).get("dp_pair_coordination", {}).get( + "all_expected_dp_peers_observed" + ) + ) + for cell in cells.values() + ), + "all_shared_dp_pair_coordination_signatures_identical": all( + bool( + cell.get("worker", {}).get("dp_pair_coordination", {}).get( + "all_shared_epoch_signatures_identical" + ) + ) for cell in cells.values() ), "all_expected_workers_observed": ( @@ -259,7 +503,7 @@ def markdown(payload: dict[str, Any]) -> str: "", "P0 verifies the collective/liveness premise only; it is not a performance result.", "", - "| policy | complete usage-verified requests | candidate K values observed | worker ranks | rank phase counts equal | within-DP execution sequences identical |", + "| policy | complete usage-verified requests | candidate K values observed | worker ranks | target phase counts equal within each DP replica | target sequences identical within each DP replica |", "|---|---:|---:|---:|---:|---:|", ] for name, cell in payload["cells"].items(): @@ -275,8 +519,8 @@ def markdown(payload: dict[str, Any]) -> str: valid=probe.get("valid", False), values=values, ranks=worker.get("worker_rank_count", 0), - equal=worker.get("phase_count", {}).get("all_equal", False), - within=worker.get("all_within_dp_sequences_identical", False), + equal=worker.get("all_target_phase_counts_equal_within_dp_replica", False), + within=worker.get("all_target_sequences_identical_within_dp_replica", False), ) ) lines.extend( @@ -319,7 +563,7 @@ def main() -> int: args.output_json.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") args.output_md.write_text(markdown(payload), encoding="utf-8") print(json.dumps(payload, indent=2, sort_keys=True)) - return 0 if payload["data_sanity"]["invariants"]["all_probe_integrity_valid"] else 2 + return 0 if all(payload["data_sanity"]["invariants"].values()) else 2 if __name__ == "__main__":