Repository navigation
Fix gemm_multiple_run to report the slowest TPU device - #140
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a report_slowest_device option to TPU benchmarking utilities, allowing the slowest device's durations to be reported when XLA:TPU lacks an end-of-program barrier. The feedback suggests several improvements: safely retrieving the pid from trace events to avoid potential KeyErrors, removing a verbose debug print statement that clutters logs, and addressing a potential bypass of the new logic when no marker events are found.
54f0bc1 to
1341174
Compare
1341174 to
9009729
Compare
| run_on_local_node: bool = False, | ||
| report_slowest_device: bool = True, | ||
| ) -> Dict[str, Any]: | ||
| del report_slowest_device |
There was a problem hiding this comment.
Because run_benchmark.py passes the exact same YAML config dictionary (**benchmark_param) to both functions.
There was a problem hiding this comment.
Rather than del, we should consider recording report_slowest_device in metadata such that it's clear if we're using the designated mode.
Earlier, per-iteration durations were taken only from the lowest-PID device (TPU-0). XLA:TPU has no end-of-program barrier, so a slower non-zero device does not show up in TPU-0's timings, and a degraded chip on the host went unreported. Now gemm_multiple_run groups trace durations by device and reports the device with the lowest median throughput. Median is used because CHS compares the reported Throughput (median) against its threshold. That figure is the median of total_flops / duration, and total_flops is the same on every device, so selecting by median of 1/duration picks exactly the device with the lowest throughput CHS checks. Testing - Unit tests on synthetic traces: slowest device selected across 2 and 4 devices, median-of-1/duration rule with an even number of runs, events without device_duration_ps ignored, default behavior still returns TPU-0, and gemm_multiple_run passes report_slowest_device=True.
9009729 to
1fc92f2
Compare
2 participants
Earlier, per-iteration durations were taken only from the lowest-PID device (TPU-0). XLA:TPU has no end-of-program barrier, so a slower non-zero device does not show up in TPU-0's timings, and a degraded chip on the host went unreported.
Now gemm_multiple_run groups trace durations by device and reports the device with the lowest median throughput.
Median is used because CHS compares the reported Throughput (median) against its threshold. That figure is the median of total_flops / duration, and total_flops is the same on every device, so selecting by median of 1/duration picks exactly the device with the lowest throughput CHS checks.
Testing