compute_standardization_stats.py --distributed silently corrupts saved stats via positional depadding after gather

[nikhil3495]

Description

compute_standardization_stats.py's --distributed path pads WeatherDataset
(via PaddedWeatherDataset) so its length divides evenly across ranks, then
after the multi-rank gather "depads" by taking a positional prefix:

means, squares = (
    [means_gathered_tensor[i] for i in original_indices],
    [squares_gathered_tensor[i] for i in original_indices],
)

(and equivalently diff_means_gathered_tensor[:n_original_windows] for the
diff stats), where original_indices = list(range(total_samples)). This
assumes the padded rows land at the tail of gathered_tensor.

They don't. DistributedSampler(shuffle=False) assigns this rank the dataset
indices range(len(dataset))[rank::world_size] — a stride of world_size,
not a contiguous block. So whenever total_samples % world_size != 0, the
padded indices (>= total_samples) land on different ranks depending on
padded_index % world_size, not all on the same rank. After
all_gather_object concatenates the per-rank results rank-major, the
padding is scattered through the middle of the tensor, not sitting at the
tail.

Taking the first total_samples entries then silently keeps a few padded
(duplicated last-sample) rows from one rank while dropping an equal number
of real rows from another — corrupting the saved parameter_mean.pt,
parameter_std.pt, diff_mean.pt and diff_std.pt used to standardize the
whole dataset for training.

Reproduction

Using this repo's actual torch.utils.data.distributed.DistributedSampler
with total_samples=101, world_size=4, batch_size=8
(padded_samples=3, padded length 104):

from torch.utils.data.distributed import DistributedSampler

class DS:
    def __init__(self, n): self.n = n
    def __len__(self): return self.n

total_samples, world_size = 101, 4
padded_len = 104  # PaddedWeatherDataset's computed length
ds = DS(padded_len)

per_rank = [list(DistributedSampler(ds, num_replicas=world_size, rank=r, shuffle=False))
            for r in range(world_size)]
flattened = [idx for rank_list in per_rank for idx in rank_list]

naive_selected = flattened[:total_samples]
print("padded rows wrongly INCLUDED:", [i for i in naive_selected if i >= total_samples])
print("real rows wrongly DROPPED:", [i for i in range(total_samples) if i not in naive_selected])

Output:

padded rows wrongly INCLUDED: [101, 102]
real rows wrongly DROPPED: [95, 99]

Not covered by existing tests — tests/test_compute_standardization_stats.py
(added in #411) only checks that slicing preserves tensor shape on
synthetic data; it never simulates an actual multi-rank
DistributedSampler, which is why this gap went uncaught by that PR's
otherwise-thorough fix for the neighboring shape/IndexError bugs
(#409/#412/#413).

Fix

Identify real vs. padded rows locally, per rank, before the gather —
using the sampler's own (already deterministic) per-rank index order — and
filter each minibatch's contribution as soon as it's computed, instead of
trying to reconstruct sample identity positionally after the gather. Order
doesn't matter for a mean/std reduction, so no reordering is needed once
padding is dropped at the source.