MIR.utils.distributed¶

Multi-GPU helpers for training and inference with DistributedDataParallel or DataParallel.

Launch with torchrun --nproc_per_node=N train.py and call init_distributed() once at start-up. The helpers also work unchanged in a single process (one GPU or CPU), so the same script runs everywhere.

Checkpoints: wrapped models prefix their state_dict keys with module.. Save with state_dict_for_saving() (or unwrap_model(model).state_dict()) so the file loads into a plain model with strict=True, exactly like the published MIR weights; strip_module_prefix() fixes checkpoints that were saved from a wrapped model.

Example:

from MIR.utils import distributed as dist_utils

device = dist_utils.init_distributed()                 # torchrun env vars, or a single process
model = dist_utils.wrap_model(model.to(device), device=device)
sampler = dist_utils.distributed_sampler(train_set)    # None in a single process
loader = DataLoader(train_set, batch_size=1, sampler=sampler, shuffle=sampler is None)
for epoch in range(n_epochs):
    if sampler is not None:
        sampler.set_epoch(epoch)
    ...
    loss_value = dist_utils.all_reduce_mean(loss.detach())
if dist_utils.is_main_process():
    torch.save({"state_dict": dist_utils.state_dict_for_saving(model)}, path)

Functions

`all_reduce_mean`(value)	Average a scalar or tensor over all processes (returned unchanged when not distributed).
`barrier`()	Synchronise all processes (no-op when not distributed).
`cleanup`()	Destroy the process group (no-op when not distributed).
`distributed_sampler`(dataset[, shuffle, ...])	`DistributedSampler` for `dataset` when distributed, otherwise `None`.
`get_rank`()	Rank of this process (0 when not distributed).
`get_world_size`()	Number of processes (1 when not distributed).
`init_distributed`([backend, device_type])	Initialise the default process group from the `torchrun` environment and return this process's device.
`is_distributed`()	True when a `torch.distributed` process group is initialised with more than one process.
`is_main_process`()	True on rank 0 (and in a single process); use it to guard logging and checkpoint saving.
`state_dict_for_saving`(model)	`state_dict` without wrapper prefixes, loadable into the plain model with `strict=True`.
`strip_module_prefix`(state_dict[, prefix])	Remove a leading `"module."` (as written by wrapped models) from every key of `state_dict`.
`unwrap_model`(model)	Return the underlying module of a `DataParallel` / `DistributedDataParallel` wrapper (or `model`).
`wrap_model`(model[, mode, device, ...])	Wrap a model for multi-GPU training or inference.