credit.distributed

credit.distributed#

Attributes#

logger

DISTRIBUTED_MODES

DEFAULT_MASTER_PORT

Functions#

select_device(→ torch.device)

Pick the torch device with precedence cuda → mps → cpu.

setup(rank, world_size, mode[, backend, device_id])

Initializes the distributed process group.

resolve_master_addr()

Resolve a routable (non-loopback) IP address for the current host.

resolve_socket_ifname()

Return the network interface name owning this host's routable address.

configure_gloo_ifname()

Bind Gloo to a routable interface when GLOO_SOCKET_IFNAME is unset.

get_rank_info(trainer_mode)

Gets rank and size information for distributed training.

should_not_checkpoint(module)

distributed_model_wrapper(conf, neural_network, device)

Wraps the neural network model for distributed training.

distributed_model_wrapper_gen2(conf, model, device)

Wrap a model for V2 distributed training.

Module Contents#

credit.distributed.logger#
credit.distributed.DISTRIBUTED_MODES = ('fsdp', 'fsdp2', 'ddp', 'domain_parallel', 'fsdp+domain_parallel')#
credit.distributed.select_device(local_rank: int = 0, device: str | torch.device | None = None, set_benchmark: bool = False) → torch.device#

Pick the torch device with precedence cuda → mps → cpu.

Centralizes device selection so every application entry point agrees on MPS support for Apple Silicon Macs (previously only the rollout apps had an MPS branch; training and preprocessing fell back to CPU).

Parameters:

  • local_rank – Local rank for multi-GPU device assignment. Ignored when device is provided or when CUDA is unavailable.

  • device – Optional explicit device override (e.g. from a CLI --device arg). When provided it takes precedence over auto-detection.

  • set_benchmark – When True, enable cudnn.benchmark if a CUDA device is selected. Off by default because trainers call this after seed_everything has set cudnn.benchmark = False for determinism — flipping it back on would silently re-enable nondeterministic autotuning.

Returns:

torch.device — and as a side effect calls torch.cuda.set_device when a CUDA device (auto-detected or explicit with an index) is selected.

credit.distributed.setup(rank, world_size, mode, backend='nccl', device_id=None)#

Initializes the distributed process group.

Parameters:

  • rank (int) – The rank of the process within the distributed setup.

  • world_size (int) – The total number of processes in the distributed setup.

  • mode (str) – The mode of operation (e.g., ‘fsdp’, ‘ddp’).

  • backend (str, optional) – The backend to use for distributed training. Defaults to ‘nccl’.

  • device_id (torch.device, optional) – Local CUDA device. Passed to init_process_group to suppress the PyTorch 2.3+ barrier() UserWarning about missing device_id.

credit.distributed.DEFAULT_MASTER_PORT = 29500#
credit.distributed.resolve_master_addr()#

Resolve a routable (non-loopback) IP address for the current host.

socket.gethostbyname(socket.gethostname()) frequently returns a loopback address (127.0.0.1) on HPC nodes whose hostname maps to localhost in /etc/hosts, which makes it unusable as a rendezvous address. This prefers a non-loopback result from hostname resolution, and otherwise discovers the outbound-interface IP by opening a dummy UDP socket (no packets are actually sent).

Returns:

A best-effort non-loopback IPv4 address, falling back to

127.0.0.1 if none can be determined.

Return type:

str

credit.distributed.resolve_socket_ifname()#

Return the network interface name owning this host’s routable address.

Used to set GLOO_SOCKET_IFNAME so Gloo binds to a real (non-loopback) interface instead of falling back to 127.0.0.1.

Returns:

The matching interface name, or None if it cannot be

determined (no routable address, psutil unavailable, or no matching interface), in which case the caller should leave GLOO_SOCKET_IFNAME unset.

Return type:

str | None

credit.distributed.configure_gloo_ifname()#

Bind Gloo to a routable interface when GLOO_SOCKET_IFNAME is unset.

PyTorch’s Gloo backend (used for CPU-side collectives even in NCCL runs) binds to loopback when it cannot resolve the hostname to a local address, which breaks those collectives across nodes. This sets GLOO_SOCKET_IFNAME to the interface owning this host’s routable address as a best-effort default. It is a no-op if the variable is already set or no interface can be determined, and it intentionally leaves NCCL_SOCKET_IFNAME alone so NCCL is free to auto-select the high-speed interconnect.

credit.distributed.get_rank_info(trainer_mode)#

Gets rank and size information for distributed training.

Parameters:

trainer_mode (str) – The mode of training (e.g., ‘fsdp’, ‘ddp’). Anything outside DISTRIBUTED_MODES is treated as single-process.

Returns:

A tuple containing LOCAL_RANK (int), WORLD_RANK (int), and WORLD_SIZE (int).

Return type:

tuple

credit.distributed.should_not_checkpoint(module)#
credit.distributed.distributed_model_wrapper(conf, neural_network, device)#

Wraps the neural network model for distributed training.

Supports modes: ‘fsdp’, ‘ddp’, ‘domain_parallel’, ‘fsdp+domain_parallel’.

For domain_parallel modes, the model’s Conv2d/Conv3d/ConvTranspose2d/GroupNorm layers are replaced with domain-parallel equivalents that handle halo exchange and distributed normalization. For fsdp+domain_parallel, domain-parallel conversion is done first, then FSDP wrapping uses the data-parallel subgroup.

Parameters:

  • conf (dict) – The configuration dictionary containing training settings.

  • neural_network (torch.nn.Module) – The neural network model to be wrapped.

  • device (torch.device) – The device on which the model will be trained.

Returns:

The wrapped model ready for distributed training.

Return type:

torch.nn.Module

credit.distributed.distributed_model_wrapper_gen2(conf: dict, model, device)#

Wrap a model for V2 distributed training.

Reads trainer.parallelism and applies, in order:

  1. Domain parallelism (spatial H-shard, Negin’s layer swap)

  2. Tensor parallelism (Col/Row linear split across TP group)

  3. FSDP2 or DDP over the data-parallel group

V1 trainer code still calls distributed_model_wrapper — this is V2 only.

Parameters:

  • conf – Full training config.

  • model – Raw nn.Module.

  • device – torch.device for this rank.

Returns:

Wrapped model (and stores domain_manager on it if domain > 1).