mapping_networks.distributed.wrap_ddp

mapping_networks.distributed.wrap_ddp(model: MappingModel, device_ids: list[int] | None = None, output_device: int | None = None, find_unused_parameters: bool = False, **kwargs: Any) → torch.nn.parallel.DistributedDataParallel

Wrap a MappingModel for DDP training.

Synchronization scope: DDP synchronizes gradients for all nn.Parameter tensors in the wrapped module. In a MappingModel:

  • Trainable latent vectors are nn.Parameter objects and will be synchronized — this is correct and desired.

  • Frozen mapper weights are registered as buffers (requires_grad=False) and are not synchronized — they are byte-for-byte identical on every rank because they share the same random seed and initialization, so no communication is needed.

  • The frozen target model’s parameters are also registered as buffers and excluded from synchronization for the same reason.

Parameters:

  • model – A MappingModel instance that has already been moved to the target device (e.g. model.cuda()).

  • device_ids – List of GPU device IDs for this process. Pass None for CPU-only (gloo) training.

  • output_device – Device where output will be gathered. Defaults to device_ids[0] when device_ids is set.

  • find_unused_parameters – Set to True only if your generator may skip some latent vectors during certain forward passes. The default False avoids the overhead of the unused-parameter scan.

  • **kwargs – Additional keyword arguments forwarded to torch.nn.parallel.DistributedDataParallel.

Returns:

A DistributedDataParallel-wrapped model. Access the underlying MappingModel via ddp_model.module.

Raises:

RuntimeError – If the process group has not been initialized.

Example:

setup_ddp(rank=0, world_size=1)
model = MappingModel(nn.Linear(4, 2), latent_dim=8).cpu()
ddp_model = wrap_ddp(model)