credit.domain_parallel.layers

credit.domain_parallel.layers#

Domain-parallel layer wrappers.

Each wrapper replaces a standard PyTorch layer with a version that handles halo exchange and/or distributed reductions for domain-parallel training.

Layers that are purely local (1x1 convolutions, channel-wise normalization, activations, etc.) do not need wrappers and are left unchanged.

Classes#

`DomainParallelConv2d`	Domain-parallel Conv2d with automatic halo exchange.
`DomainParallelCrossEmbedBranch`	Domain-parallel wrapper for one CrossEmbedLayer kernel branch.
`DomainParallelConv3d`	Domain-parallel Conv3d with halo exchange along the lat dimension.
`DomainParallelConvTranspose2d`	Domain-parallel ConvTranspose2d with halo exchange.
`DomainParallelConvTranspose3d`	Domain-parallel ConvTranspose3d with halo exchange along the lat dimension.
`DomainParallelGroupNorm`	Domain-parallel GroupNorm with distributed statistics.
`DomainParallelPeriodicConv2d`	Domain-parallel wrapper for PeriodicConv2d.
`DomainParallelInterpolate`	Domain-parallel bilinear interpolation.

Module Contents#

class credit.domain_parallel.layers.DomainParallelConv2d(conv, shard_dim=-2)#

Bases: torch.nn.Module

Domain-parallel Conv2d with automatic halo exchange.

Before the convolution, exchanges halo rows along the sharding dimension (latitude by default) so boundary pixels see correct neighbors. After the convolution, the output has the correct local shard size.

For strided convolutions (like in CrossEmbedLayer), the halo width is (kernel_size - stride) // 2 to account for the stride’s effect on output boundaries.

Parameters:

conv – An existing nn.Conv2d module to wrap.
shard_dim – Spatial dimension being sharded (-2 for H in BCHW).

conv#

shard_dim = -2#

halo_exchange#

forward(x)#

property weight#

property bias#

class credit.domain_parallel.layers.DomainParallelCrossEmbedBranch(branch, shard_dim=-2)#

Bases: torch.nn.Module

Domain-parallel wrapper for one CrossEmbedLayer kernel branch.

CrossEmbedLayer (credit/models/wxformer/crossformer.py) pads each kernel branch explicitly with nn.ZeroPad2d before an otherwise-unpadded Conv2d, so every kernel size in the multi-scale conv gets the exact “same” output shape regardless of (kernel - stride) parity. Domain-parallel conversion must NOT apply that padding along the sharded (H) dimension as-is: at an interior domain-shard boundary the correct “padding” is real neighbor data from the adjacent shard (halo exchange), not zeros – zero-padding an interior boundary is wrong, and it also corrupts the halo exchange, which then exchanges those fabricated zeros as if they were real boundary data. Along the un-sharded (W) dimension the zero-padding is still exactly right and must be preserved.

Plain DomainParallelConv2d isn’t sufficient here: credit/domain_parallel/ convert.py’s module walk finds only the inner Conv2d inside the branch’s nn.Sequential(ZeroPad2d, Conv2d) and wraps that alone, leaving the sibling ZeroPad2d in place – it keeps firing unconditionally, before DomainParallelConv2d’s halo exchange ever sees the input. And even if the inner Conv2d were wrapped directly, DomainParallelConv2d.forward() falls back to the wrapped conv’s own .padding attribute for the W direction, which is always (0, 0) here: all of this branch’s padding was deliberately moved out of the Conv2d and into the ZeroPad2d. So the W padding must be read from the ZeroPad2d and reapplied explicitly instead.

convert.py matches this class (CrossEmbedConvBranch, a distinctly-named nn.Sequential subclass) by isinstance, ahead of the generic Conv2d rule, and replaces the whole branch with this wrapper – not just its inner Conv2d.

Parameters:

branch – A CrossEmbedConvBranch (nn.Sequential(ZeroPad2d, Conv2d)) to wrap.
shard_dim – Spatial dimension being sharded (-2 for H in BCHW).

conv#

shard_dim = -2#

w_padding#

halo_exchange#

forward(x)#

property weight#

property bias#

class credit.domain_parallel.layers.DomainParallelConv3d(conv, shard_dim=3)#

Bases: torch.nn.Module

Domain-parallel Conv3d with halo exchange along the lat dimension.

Used for CubeEmbedding’s patch-embedding Conv3d. The sharded dimension in the 5D tensor (B, C, T, H, W) is H (dim=-2 or dim=3).

Parameters:

conv – An existing nn.Conv3d module to wrap.
shard_dim – Spatial dimension being sharded (3 for H in 5D BCTHW).

conv#

shard_dim = 3#

halo_exchange#

forward(x)#

class credit.domain_parallel.layers.DomainParallelConvTranspose2d(conv, shard_dim=-2)#

Bases: torch.nn.Module

Domain-parallel ConvTranspose2d with halo exchange.

For kernel=2, stride=2 (standard upsample): no halo needed, purely local. For kernel=4, stride=2, padding=1: halo of 1 needed.

Parameters:

conv – An existing nn.ConvTranspose2d module to wrap.
shard_dim – Spatial dimension being sharded (-2 for H in BCHW).

conv#

shard_dim = -2#

halo_width#

halo_exchange#

forward(x)#

class credit.domain_parallel.layers.DomainParallelConvTranspose3d(conv, shard_dim=3)#

Bases: torch.nn.Module

Domain-parallel ConvTranspose3d with halo exchange along the lat dimension.

For kernel==stride (e.g. Pangu PatchRecovery with patch_size=(2,4,4)): no halo needed — purely local upsampling. For kernel>stride in H: halo exchange required.

Parameters:

conv – An existing nn.ConvTranspose3d module to wrap.
shard_dim – Spatial dimension being sharded (3 for H in BCZHW).

conv#

shard_dim = 3#

halo_width#

halo_exchange#

forward(x)#

class credit.domain_parallel.layers.DomainParallelGroupNorm(norm)#

Bases: torch.nn.Module

Domain-parallel GroupNorm with distributed statistics.

GroupNorm computes mean and variance over (H, W) for each group of channels. With domain-sharded H, we need to all_reduce the statistics across the domain group before applying normalization.

Parameters:: norm – An existing nn.GroupNorm module to wrap.

num_groups#

num_channels#

eps#

affine#

forward(x)#

class credit.domain_parallel.layers.DomainParallelPeriodicConv2d(periodic_conv, shard_dim=-2)#

Bases: torch.nn.Module

Domain-parallel wrapper for PeriodicConv2d.

PeriodicConv2d manually pads W with circular (longitude) and H with reflect (latitude) before calling an inner nn.Conv2d(padding=0). With domain-sharded H, reflect padding on H is wrong — boundary ranks would reflect against the shard edge instead of the true pole.

This wrapper replaces the reflect padding on H with halo exchange, while keeping the circular padding on W unchanged.

Parameters:

periodic_conv – An existing PeriodicConv2d module to wrap. Must have a .conv (nn.Conv2d) and .padding (int) attribute.
shard_dim – Spatial dimension being sharded (-2 for H in BCHW).

conv#

padding#

shard_dim = -2#

halo_exchange#

forward(x)#

property weight#

property bias#

class credit.domain_parallel.layers.DomainParallelInterpolate(size, mode='bilinear', shard_dim=-2)#

Bases: torch.nn.Module

Domain-parallel bilinear interpolation.

Exchanges a 1-row halo before interpolation so boundary pixels interpolate correctly, then trims the extra output rows.

Parameters:

size – Target output size (H, W).
mode – Interpolation mode (default: ‘bilinear’).
shard_dim – Spatial dimension being sharded (-2 for H in BCHW).

size#

mode = 'bilinear'#

shard_dim = -2#

halo_exchange#

forward(x)#