Stop trim_mean sorting the caller's list in place (#8199)
## The bug
`CommsLogger.comms_dict` stores parallel lists per message size,
`[count, latencies, algbws, busbws]`, where index `i` is the i-th
recorded op. `get_operation_summary()` and `log_all()` hand each of
those lists to `trim_mean`, which sorted **in place**:
```python
data.sort()
k = int(round(n * (trim_percent)))
return mean(data[k:n - k])
```
So summarising sorts the three lists independently and destroys the
correspondence between them. Driving the real `CommsLogger` with
latencies `[3.0, 1.0, 2.0]`:
```
BEFORE latency: [3.0, 1.0, 2.0] algbw: [0.0055, 0.0164, 0.0082]
AFTER latency: [1.0, 2.0, 3.0] algbw: [0.0055, 0.0082, 0.0164]
stored data mutated by a read-only summary call: True
latency[i] * algbw[i], constant by construction: [0.0055, 0.0164, 0.0492]
```
`algbw` is computed from `latency` in `calc_bw_log`, so `latency[i] *
algbw[i]` is constant for a fixed message size. After summarising it is
not, because row 0 now pairs the fastest op (1.0 ms) with the lowest
bandwidth. `get_raw_data()` hands that out to anyone consuming the log.
Two things make this look unintended rather than a quirk:
- `get_operation_summary()` already does `op_data =
self.comms_dict[operation_name].copy()` with the comment "Create a
snapshot to avoid concurrent modification issues". The intent not to
disturb the stored records is explicit; the copy is just shallow, so the
inner lists are shared and the sort reaches them anyway.
- A summary is a read. Calling it should not rewrite what was recorded,
and nothing documents that it does.
## It also breaks the straggler breakdown, which is the bigger effect
Raised in review and worth stating here, because it is worse than the
reordering above.
`log_all()` reads `vals[1]` twice. The summary loop calls
`trim_mean(vals[1], 0.1)`, which sorts it, and the `show_straggler`
block afterwards builds both `lats` and `min_lats` from that same,
now-sorted list:
```python
lats = torch.tensor(vals[1], device=device)
min_lats = torch.tensor(vals[1], device=device)
dist.all_reduce(min_lats, op=ReduceOp.MIN)
total_straggler = (lats - min_lats).sum().item()
```
`all_reduce(..., MIN)` is elementwise, so it assumes index `i` is the
same collective on every rank. Each rank has sorted its own list
independently by then, so index `i` is "the i-th fastest op on *that*
rank" and the reduction takes the minimum across unrelated operations.
`comms_dict_snapshot = self.comms_dict.copy()` at the top of `log_all()`
does not help, for the same reason the copy in `get_operation_summary()`
does not: it is shallow.
That is not a cosmetic reordering, it changes the number. Two ranks,
four index-aligned ops:
```
rank0 = [2.30, 1.60, 3.60, 1.29]
rank1 = [3.14, 2.46, 1.23, 3.03]
true total_straggler per rank [2.37, 3.44]
with the in-place sort [0.52, 1.59] <- what gets reported
```
Over 20000 random four-op cases, 16246 report a different total, and the
error has a direction: sorting both ranks aligns them as closely as
their values allow, so the differences shrink and the straggler effect
is systematically **under**-reported. Small cases can coincide (a
three-op example I tried happened to agree), which is part of why this
is easy to miss.
## Fix
`data = sorted(data)`. That fixes every caller at the shared function
rather than patching `log_all` and `get_operation_summary` separately,
and the trimmed mean it returns is unchanged.
## Tests
Added to `tests/unit/comm/test_comms_logger.py`:
- `test_get_operation_summary_does_not_reorder_the_stored_records`
populates `comms_dict` directly and asserts the three stored lists
survive a summary call unchanged, that `latency[i] * algbw[i]` is still
constant, and that `avg_latency_ms` is still the correct trimmed mean.
- `test_trim_mean_does_not_mutate_its_argument` pins the contract at the
function itself.
Both are dist-free, like the existing test in that file. `comms_dict` is
populated directly rather than through `append()` because `append()`
calls `calc_bw_log`, which needs a live process group.
Fail-before / pass-after against the unmodified `timer.py`:
```
upstream timer.py, new tests present
PASS test_stop_profiling_comms_disables_prof_all
FAIL test_get_operation_summary_does_not_reorder_the_stored_records
FAIL test_trim_mean_does_not_mutate_its_argument
with this fix
PASS all three
```
The pre-existing test passing either way is deliberate: this is a
distinct failure from the one it covers.
How these were run, since I would rather say than imply a normal pytest
run: I do not have a GPU or a built DeepSpeed here, so I executed
`deepspeed/utils/timer.py` and `deepspeed/utils/comms_logging.py` from
source against stub `deepspeed.comm` / `deepspeed.accelerator` modules
and ast-extracted the tests. That exercises the real `CommsLogger` and
the real `trim_mean`; CI is the runner for the suite proper.
No existing issue or PR covers this; searching `trim_mean` across open
and closed returns only the merged PR that introduced the current form.
---------
Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
