Aggregate (provider, n_hidden_states, tp) to mean/min/max in microseconds.
(long: pd.DataFrame, args: "Args")
| 178 | max_slowdown: float = 2.0 # drop runs whose FMMS median is more than this factor slower than the fastest at the same tp |
| 179 | |
| 180 | |
| 181 | def write_summary_csv(long: pd.DataFrame, args: "Args") -> Path: |
| 182 | """Aggregate (provider, n_hidden_states, tp) to mean/min/max in microseconds.""" |
| 183 | sel = long.query("n_hidden_states in @args.h_values and provider in @DEFAULT_PROVIDERS").copy() |
| 184 | sel["time[us]"] = sel["time[ms]"] * 1000 |
| 185 | summary = ( |
| 186 | sel.groupby(["provider", "n_hidden_states", "tp"])["time[us]"] |
| 187 | .agg(n_runs="count", mean_time_us="mean", min_time_us="min", max_time_us="max") |
| 188 | .round(2) |
| 189 | .reset_index() |
| 190 | .rename( |
| 191 | columns={ |
| 192 | "provider": "method", |
| 193 | "n_hidden_states": "batch_size", |
| 194 | "tp": "tensor_parallel_size", |
| 195 | } |
| 196 | ) |
| 197 | ) |
| 198 | summary["batch_size"] = summary["batch_size"].astype(int) |
| 199 | summary = summary.sort_values(["method", "batch_size", "tensor_parallel_size"]) |
| 200 | summary["range_us"] = (summary["max_time_us"] - summary["min_time_us"]).round(2) |
| 201 | csv_path = args.out.with_name("tp-scaling-runs-summary.csv") |
| 202 | summary.to_csv(csv_path, index=False) |
| 203 | return csv_path |
| 204 | |
| 205 |