115 lines
3.5 KiB
Python
115 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Run Frontier with the existing joint repair plus a whole-layer decode curve."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import runpy
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
REPO = ROOT.parents[1]
|
|
GRID = json.loads(Path(os.environ["FRONTIER_WHOLE_DECODE_GRID"]).read_text())
|
|
WHOLE_CURVE = {
|
|
str(tp): {
|
|
str(cell["batch"]): float(cell["median_execute_ms"])
|
|
for cell in GRID["cells"]
|
|
if int(cell["tp"]) == tp
|
|
}
|
|
for tp in (2, 4)
|
|
}
|
|
|
|
# Apply the existing serving-path collective/MoE correction first.
|
|
joint = runpy.run_path(
|
|
str(REPO / "runs/frontier-collective-joint-v0/run_frontier_with_curves.py")
|
|
)
|
|
USAGE = joint["USAGE"]
|
|
_pure_decode_point = joint["_pure_decode_point"]
|
|
|
|
from frontier.entities import ExecutionTime # noqa: E402
|
|
from frontier.execution_time_predictor.sklearn_moe_execution_time_predictor import ( # noqa: E402
|
|
SklearnMoEExecutionTimePredictor,
|
|
)
|
|
|
|
|
|
_JOINT_STAGE_PREDICT = SklearnMoEExecutionTimePredictor.predict_stage_execution_time
|
|
|
|
|
|
def _whole_decode_stage_time(
|
|
self,
|
|
batch,
|
|
stage_id,
|
|
cluster_type,
|
|
num_layers=1,
|
|
layer_id=0,
|
|
):
|
|
base = _JOINT_STAGE_PREDICT(
|
|
self,
|
|
batch,
|
|
stage_id,
|
|
cluster_type,
|
|
num_layers=num_layers,
|
|
layer_id=layer_id,
|
|
)
|
|
point = _pure_decode_point(self, batch)
|
|
if point is None or point[0] == "1" or point[1] == "1":
|
|
return base
|
|
tp, decode_batch = point
|
|
if tp not in WHOLE_CURVE or decode_batch not in WHOLE_CURVE[tp]:
|
|
raise ValueError(
|
|
"Whole-layer curve has no exact pure-decode point for "
|
|
f"TP={tp}, batch={decode_batch}; refusing to extrapolate"
|
|
)
|
|
layers = int(base.num_layers)
|
|
if layers <= 0:
|
|
raise ValueError(f"invalid stage layer count: {layers}")
|
|
target_ms = WHOLE_CURVE[tp][decode_batch]
|
|
USAGE[f"whole_decode:tp{tp}-b{decode_batch}:target_ms={target_ms:.9f}"] += 1
|
|
return ExecutionTime(
|
|
num_layers_per_pipeline_stage=layers,
|
|
attention_rope_execution_time=0.0,
|
|
attention_kv_cache_save_execution_time=0.0,
|
|
attention_decode_execution_time=0.0,
|
|
attention_prefill_execution_time=0.0,
|
|
attention_layer_pre_proj_execution_time=0.0,
|
|
attention_layer_post_proj_execution_time=0.0,
|
|
attn_norm_time=0.0,
|
|
mlp_norm_time=0.0,
|
|
add_time=0.0,
|
|
tensor_parallel_communication_time=0.0,
|
|
pipeline_parallel_communication_time=0.0,
|
|
expert_parallel_communication_time=0.0,
|
|
moe_gating_time=0.0,
|
|
moe_shuffling_time=0.0,
|
|
schedule_time=base._schedule_time,
|
|
sampler_e2e_time=base._sampler_e2e_time,
|
|
prepare_inputs_e2e_time=base._prepare_inputs_e2e_time,
|
|
process_model_outputs_time=base._process_model_outputs_time,
|
|
ray_comm_time=base._ray_comm_time,
|
|
is_moe=True,
|
|
moe_grouped_gemm_time=target_ms / layers,
|
|
pp_producer_send_path_runtime_time=(
|
|
base._pp_producer_send_path_runtime_time
|
|
),
|
|
pp_receiver_head_runtime_time=base._pp_receiver_head_runtime_time,
|
|
pp_prefill_consumer_active_runtime_time=(
|
|
base._pp_prefill_consumer_active_runtime_time
|
|
),
|
|
pp_stage_boundary_residual_runtime_time=(
|
|
base._pp_stage_boundary_residual_runtime_time
|
|
),
|
|
)
|
|
|
|
|
|
SklearnMoEExecutionTimePredictor.predict_stage_execution_time = (
|
|
_whole_decode_stage_time
|
|
)
|
|
|
|
from frontier.main import main # noqa: E402
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|