gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Analyze corrected Frontier op traces

Browse Source
This commit is contained in:
Gahow Wang
2026-07-19 18:53:12 +08:00
parent d5c3c93577
commit 3349b23290
2 changed files with 166 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

164
runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py
View File

@@ -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"]}))