From 3349b2329011fe708002ae04a2ddb5d84a0ead74 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Sun, 19 Jul 2026 18:53:12 +0800 Subject: [PATCH] Analyze corrected Frontier op traces --- .../analyze_qwen235_fixed_pd_state.py | 164 +++++++++++++++++- .../test_fidelity_envelope.py | 5 + 2 files changed, 166 insertions(+), 3 deletions(-) diff --git a/runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py b/runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py index 8843574..03d021d 100644 --- a/runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py +++ b/runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py @@ -87,6 +87,66 @@ CATEGORIES = { }, } +OP_CATEGORIES = { + "attention": { + "input_layernorm", + "attn_pre_proj", + "attn_rope", + "attn_prefill", + "attn_decode", + "attn_kv_cache_save", + "attn_post_proj", + }, + "dense_mlp_compute": { + "post_attention_layernorm", + "mlp_up_proj", + "mlp_act", + "mlp_down_proj", + }, + "moe_compute": { + "moe_grouped_gemm", + "share_expert_up_proj", + "share_expert_act", + "share_expert_down_proj", + }, + "moe_routing": { + "moe_gating_linear", + "moe_gating_routing_topk", + "moe_shuffling", + }, + "ep_communication": { + "expert_parallel_alltoall", + "expert_parallel_alltoall_dispatch", + "expert_parallel_alltoall_combine", + "expert_parallel_allreduce", + "expert_parallel_allreduce_wait", + }, + "tp_dp_communication": { + "attn_tensor_parallel_allreduce", + "mlp_tensor_parallel_allreduce", + "moe_tensor_parallel_allgather", + "moe_tensor_parallel_allreduce", + "share_expert_tensor_parallel_allreduce", + "dp_input_allreduce", + "dp_output_allreduce", + }, + "pipeline_communication": {"pipeline_parallel_send_recv"}, + "runtime_overhead": { + "add_attn_residual", + "add_ffn_residual", + "schedule", + "prepare_inputs_e2e", + "pp_receiver_head_runtime", + "pp_prefill_consumer_active_runtime", + "decode_draft_proposer", + "mtp_terminal_overshoot", + "pp_stage_boundary_handoff", + "sampler_e2e", + "process_model_outputs", + "ray_comm_time", + }, +} + def atomic_write(path: Path, text: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) @@ -183,6 +243,13 @@ def categorized_components(components: dict[str, Any]) -> dict[str, float]: } +def op_category(name: str) -> str: + matches = [category for category, names in OP_CATEGORIES.items() if name in names] + if len(matches) != 1: + raise ValueError(f"unclassified or multiply classified op trace event: {name}") + return matches[0] + + def load_stage_rows(state_root: Path, config: str) -> tuple[list[dict[str, Any]], dict]: result_path = state_root / config / "result.json" result = json.loads(result_path.read_text()) @@ -218,6 +285,82 @@ def load_stage_rows(state_root: Path, config: str) -> tuple[list[dict[str, Any]] return rows, result +def load_op_trace_rows( + op_trace_root: Path, config: str +) -> tuple[list[dict[str, Any]], dict, dict[str, Any]]: + result_path = op_trace_root / config / "result.json" + result = json.loads(result_path.read_text()) + if result.get("status") != "PASS" or not result.get("op_trace_enabled"): + raise ValueError(f"op-trace replay did not pass: {result_path}") + trace_path = Path(result["state_artifacts"]["op_trace"]["path"]) + grouped: dict[int, dict[str, Any]] = {} + event_names = Counter() + with trace_path.open() as source: + for line_number, line in enumerate(source, start=1): + event = json.loads(line) + if "meta" in event and len(event) == 1: + continue + name = str(event["name"]) + category = op_category(name) + event_names[name] += 1 + batch_id = int(event["batch_id"]) + tokens = [int(value) for value in event["meta"]["num_tokens"]] + entry = grouped.setdefault( + batch_id, + { + "tokens": tokens, + "categories_ms": {key: 0.0 for key in CATEGORIES}, + "events": 0, + }, + ) + if entry["tokens"] != tokens: + raise ValueError( + f"op trace token metadata changed within batch {batch_id}" + ) + entry["categories_ms"][category] += float(event["duration_ms"]) + entry["events"] += 1 + if not grouped: + raise ValueError(f"no op trace events: {trace_path}") + + rows = [] + for batch_id, entry in sorted(grouped.items()): + tokens = entry["tokens"] + total = sum(entry["categories_ms"].values()) + rows.append( + { + "batch_id": batch_id, + "batch_size": len(tokens), + "decode_requests": sum(value == 1 for value in tokens), + "prefill_requests": sum(value > 1 for value in tokens), + "prefill_tokens": sum(value for value in tokens if value > 1), + "total_time_ms": total, + "categories_ms": entry["categories_ms"], + "per_expert_tokens": None, + } + ) + + ledger_path = Path(result["state_artifacts"]["ledger"]["path"]) + stage_span_by_batch = { + int(row["batch_id"]): + (float(row["stage_end_ts"]) - float(row["stage_start_ts"])) * 1000 + for row in load_jsonl(ledger_path) + } + ratios = [ + row["total_time_ms"] / stage_span_by_batch[row["batch_id"]] + for row in rows + if row["batch_id"] in stage_span_by_batch + and stage_span_by_batch[row["batch_id"]] > 0 + ] + validation = { + "trace_path": str(trace_path), + "trace_batches": len(rows), + "ledger_batches": len(stage_span_by_batch), + "trace_total_over_stage_span": numeric(ratios), + "event_names": dict(sorted(event_names.items())), + } + return rows, result, validation + + def weighted_component_mean( rows: Iterable[dict[str, Any]], *, weight_name: str = "batch_size" ) -> dict[str, float]: @@ -290,12 +433,20 @@ def subtract(right: dict[str, float], left: dict[str, float]) -> dict[str, float return {name: right[name] - left[name] for name in left} -def analyze(real_root: Path, state_root: Path) -> dict[str, Any]: +def analyze( + real_root: Path, state_root: Path, op_trace_root: Path | None = None +) -> dict[str, Any]: real = summarize_real(real_root) sim = {} by_batch = {} for config in CONFIGS: - rows, replay = load_stage_rows(state_root, config) + component_source = "stage_batch_ledger" + trace_validation = None + if op_trace_root is not None and (op_trace_root / config / "result.json").is_file(): + rows, replay, trace_validation = load_op_trace_rows(op_trace_root, config) + component_source = "op_trace_execution_time_override" + else: + rows, replay = load_stage_rows(state_root, config) decode_bearing = [row for row in rows if row["decode_requests"] > 0] pure_decode = [row for row in decode_bearing if row["prefill_requests"] == 0] mixed = [row for row in decode_bearing if row["prefill_requests"] > 0] @@ -325,6 +476,8 @@ def analyze(real_root: Path, state_root: Path) -> dict[str, Any]: str(batch): value for batch, value in grouped.items() }, "scorer_equivalence": replay["scorer_equivalence"], + "component_source": component_source, + "trace_validation": trace_validation, } real_reweighted = { @@ -472,6 +625,7 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--real-root", type=Path, required=True) parser.add_argument("--state-root", type=Path, required=True) + parser.add_argument("--op-trace-root", type=Path) parser.add_argument("--json-output", type=Path, required=True) parser.add_argument("--markdown-output", type=Path, required=True) return parser.parse_args() @@ -479,7 +633,11 @@ def parse_args() -> argparse.Namespace: def main() -> None: args = parse_args() - result = analyze(args.real_root.resolve(), args.state_root.resolve()) + result = analyze( + args.real_root.resolve(), + args.state_root.resolve(), + args.op_trace_root.resolve() if args.op_trace_root else None, + ) atomic_write(args.json_output, json.dumps(result, indent=2, sort_keys=True) + "\n") atomic_write(args.markdown_output, markdown(result)) print(json.dumps({"status": result["status"], "verdict": result["verdict"]})) diff --git a/runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py b/runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py index 4a784c0..7abf5d7 100644 --- a/runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py +++ b/runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py @@ -230,6 +230,11 @@ class FidelityEnvelopeTest(unittest.TestCase): components["unknown_graph_overhead"] = 1.0 with self.assertRaisesRegex(ValueError, "component schema drift"): module.categorized_components(components) + for category, names in module.OP_CATEGORIES.items(): + for name in names: + self.assertEqual(module.op_category(name), category) + with self.assertRaisesRegex(ValueError, "unclassified"): + module.op_category("unknown_graph_overhead") def test_materialize_qwen235_allreduce_requires_serving_contract(self) -> None: module = load("materialize_qwen235_v020_allreduce.py")