Support long code traces in Frontier replay
This commit is contained in:
192
runs/frontier-collective-joint-v0/run_frontier_with_curves.py
Normal file
192
runs/frontier-collective-joint-v0/run_frontier_with_curves.py
Normal file
@@ -0,0 +1,192 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run Frontier with pure-decode serving-path MoE and/or collective curves."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
csv.field_size_limit(16 * 1024 * 1024)
|
||||
|
||||
|
||||
COLLECTIVE_PAYLOAD = json.loads(
|
||||
Path(os.environ["FRONTIER_COLLECTIVE_CURVE"]).read_text()
|
||||
)
|
||||
COLLECTIVE_VARIANT = os.environ.get("FRONTIER_COLLECTIVE_CURVE_VARIANT", "drop_mean")
|
||||
if "curve_variants_ms_per_step" in COLLECTIVE_PAYLOAD:
|
||||
try:
|
||||
COLLECTIVE_CURVE = COLLECTIVE_PAYLOAD["curve_variants_ms_per_step"][
|
||||
COLLECTIVE_VARIANT
|
||||
]
|
||||
except KeyError as error:
|
||||
raise ValueError(
|
||||
f"collective curve has no variant {COLLECTIVE_VARIANT!r}"
|
||||
) from error
|
||||
else:
|
||||
if COLLECTIVE_VARIANT != "drop_mean":
|
||||
raise ValueError("legacy collective curve only supports drop_mean")
|
||||
COLLECTIVE_CURVE = COLLECTIVE_PAYLOAD.get(
|
||||
"curve_ms_per_step", COLLECTIVE_PAYLOAD
|
||||
)
|
||||
MOE_CURVE_PATH = os.environ.get("FRONTIER_FUSED_MOE_CURVE")
|
||||
MOE_CURVE = (
|
||||
json.loads(Path(MOE_CURVE_PATH).read_text()) if MOE_CURVE_PATH else None
|
||||
)
|
||||
USAGE_PATH = Path(os.environ["FRONTIER_CURVE_USAGE"])
|
||||
USAGE: Counter[str] = Counter()
|
||||
|
||||
|
||||
def write_usage() -> None:
|
||||
USAGE_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
USAGE_PATH.write_text(json.dumps(dict(sorted(USAGE.items())), indent=2) + "\n")
|
||||
|
||||
|
||||
atexit.register(write_usage)
|
||||
|
||||
from frontier.execution_time_predictor.sklearn_moe_execution_time_predictor import ( # noqa: E402
|
||||
SklearnMoEExecutionTimePredictor,
|
||||
)
|
||||
|
||||
|
||||
_ORIGINAL_ATTN_COLLECTIVE = (
|
||||
SklearnMoEExecutionTimePredictor._get_tensor_parallel_communication_time
|
||||
)
|
||||
_ORIGINAL_MOE_COLLECTIVE = (
|
||||
SklearnMoEExecutionTimePredictor._get_moe_tensor_parallel_allreduce_time
|
||||
)
|
||||
_ORIGINAL_GROUPED_GEMM = SklearnMoEExecutionTimePredictor._get_grouped_gemm_time
|
||||
_ORIGINAL_ATTN_NORM = (
|
||||
SklearnMoEExecutionTimePredictor._get_attn_norm_layer_act_execution_time
|
||||
)
|
||||
_ORIGINAL_MLP_NORM = (
|
||||
SklearnMoEExecutionTimePredictor._get_mlp_norm_layer_act_execution_time
|
||||
)
|
||||
|
||||
|
||||
def _pure_decode_point(self, batch) -> tuple[str, str] | None:
|
||||
if batch is None or not bool(getattr(batch, "is_pure_decode_batch", False)):
|
||||
return None
|
||||
tp = str(int(self._replica_config.moe_tensor_parallel_size))
|
||||
attn_tp = str(int(self._replica_config.attn_tensor_parallel_size))
|
||||
if tp != attn_tp:
|
||||
raise ValueError(
|
||||
"Serving collective curve requires equal attention/MoE TP, got "
|
||||
f"attn_tp={attn_tp}, moe_tp={tp}"
|
||||
)
|
||||
return tp, str(len(batch.requests))
|
||||
|
||||
|
||||
def _collective_path_time(self, batch, path: str) -> float:
|
||||
point = _pure_decode_point(self, batch)
|
||||
if point is None:
|
||||
original = (
|
||||
_ORIGINAL_ATTN_COLLECTIVE
|
||||
if path == "attention"
|
||||
else _ORIGINAL_MOE_COLLECTIVE
|
||||
)
|
||||
return original(self, batch)
|
||||
tp, decode_batch = point
|
||||
if tp == "1":
|
||||
original = (
|
||||
_ORIGINAL_ATTN_COLLECTIVE
|
||||
if path == "attention"
|
||||
else _ORIGINAL_MOE_COLLECTIVE
|
||||
)
|
||||
value = original(self, batch)
|
||||
if value != 0.0:
|
||||
raise ValueError(f"TP1 {path} collective must be zero, got {value}")
|
||||
USAGE[f"collective:{path}:tp1-b{decode_batch}:structural-zero"] += 1
|
||||
return value
|
||||
if tp not in COLLECTIVE_CURVE or decode_batch not in COLLECTIVE_CURVE[tp]:
|
||||
raise ValueError(
|
||||
"Collective curve has no exact pure-decode point for "
|
||||
f"TP={tp}, batch={decode_batch}; refusing to extrapolate"
|
||||
)
|
||||
layers = int(self._num_layers_per_pipeline_stage)
|
||||
if layers <= 0:
|
||||
raise ValueError(f"invalid layers per pipeline stage: {layers}")
|
||||
USAGE[f"collective:{path}:tp{tp}-b{decode_batch}"] += 1
|
||||
# The serving curve is the sum of attention and MoE output reductions over
|
||||
# the full model step. ExecutionTime later multiplies each per-layer path.
|
||||
return float(COLLECTIVE_CURVE[tp][decode_batch]) / (2 * layers)
|
||||
|
||||
|
||||
def _attention_collective_from_curve(self, batch) -> float:
|
||||
return _collective_path_time(self, batch, "attention")
|
||||
|
||||
|
||||
def _moe_collective_from_curve(self, batch) -> float:
|
||||
return _collective_path_time(self, batch, "moe")
|
||||
|
||||
|
||||
def _grouped_gemm_from_curve(self, num_tokens_or_allocation, batch=None) -> float:
|
||||
if MOE_CURVE is None:
|
||||
return _ORIGINAL_GROUPED_GEMM(
|
||||
self, num_tokens_or_allocation, batch=batch
|
||||
)
|
||||
point = _pure_decode_point(self, batch)
|
||||
if point is None:
|
||||
return _ORIGINAL_GROUPED_GEMM(
|
||||
self, num_tokens_or_allocation, batch=batch
|
||||
)
|
||||
tp, decode_batch = point
|
||||
if tp not in MOE_CURVE or decode_batch not in MOE_CURVE[tp]:
|
||||
raise ValueError(
|
||||
"Fused MoE curve has no exact pure-decode point for "
|
||||
f"TP={tp}, batch={decode_batch}; refusing to extrapolate"
|
||||
)
|
||||
layers = int(self._num_layers_per_pipeline_stage)
|
||||
if layers <= 0:
|
||||
raise ValueError(f"invalid layers per pipeline stage: {layers}")
|
||||
USAGE[f"moe:tp{tp}-b{decode_batch}"] += 1
|
||||
return float(MOE_CURVE[tp][decode_batch]) / layers
|
||||
|
||||
|
||||
def _norm_without_fused_duplicate(self, batch, *, name: str) -> float:
|
||||
original = _ORIGINAL_ATTN_NORM if name == "attn" else _ORIGINAL_MLP_NORM
|
||||
value = original(self, batch)
|
||||
point = _pure_decode_point(self, batch)
|
||||
if point is None or point[0] == "1":
|
||||
return value
|
||||
if value < 0:
|
||||
raise ValueError(f"negative {name} norm prediction cannot be deducted: {value}")
|
||||
tp, decode_batch = point
|
||||
USAGE[f"fused_norm_deduction:{name}:tp{tp}-b{decode_batch}"] += 1
|
||||
# AllReduceFusionPattern 1 already contains residual add + RMSNorm. Returning
|
||||
# zero removes the exact predictor value that would otherwise be emitted in
|
||||
# the frozen Frontier attn_norm_time/mlp_norm_time ledger row.
|
||||
return 0.0
|
||||
|
||||
|
||||
def _attn_norm_without_fused_duplicate(self, batch) -> float:
|
||||
return _norm_without_fused_duplicate(self, batch, name="attn")
|
||||
|
||||
|
||||
def _mlp_norm_without_fused_duplicate(self, batch) -> float:
|
||||
return _norm_without_fused_duplicate(self, batch, name="mlp")
|
||||
|
||||
|
||||
SklearnMoEExecutionTimePredictor._get_tensor_parallel_communication_time = (
|
||||
_attention_collective_from_curve
|
||||
)
|
||||
SklearnMoEExecutionTimePredictor._get_moe_tensor_parallel_allreduce_time = (
|
||||
_moe_collective_from_curve
|
||||
)
|
||||
SklearnMoEExecutionTimePredictor._get_grouped_gemm_time = _grouped_gemm_from_curve
|
||||
SklearnMoEExecutionTimePredictor._get_attn_norm_layer_act_execution_time = (
|
||||
_attn_norm_without_fused_duplicate
|
||||
)
|
||||
SklearnMoEExecutionTimePredictor._get_mlp_norm_layer_act_execution_time = (
|
||||
_mlp_norm_without_fused_duplicate
|
||||
)
|
||||
|
||||
from frontier.main import main # noqa: E402
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user