DDP Training Pipeline

NeuralPOM uses Distributed Data Parallel (DDP) training across multiple GPUs with the NCCL backend. The training pipeline is implemented in the NPOMTrainerDDP class, which orchestrates the full physics-in-the-loop training loop including core pool management, multi-step physics integration, loss computation, and gradient updates.

Distributed Setup

Training is launched with torchrun (or equivalent) and initializes the distributed environment:

self.local_rank = int(os.environ.get("LOCAL_RANK", 0))
self.global_rank = int(os.environ.get("RANK", 0))
self.world_size = int(os.environ.get("WORLD_SIZE", 1))

self.device = torch.device(f"cuda:{self.local_rank}")
torch.cuda.set_device(self.device)
dist.init_process_group(backend="nccl")

Data loading uses DistributedSampler to partition the dataset across GPUs without overlap:

self.sampler = DistributedSampler(train_ds, shuffle=True)
self.train_loader = DataLoader(
    train_ds,
    batch_size=args.batch_size,
    sampler=self.sampler,
    num_workers=0,
    pin_memory=False,
    drop_last=True,
)

The sampler's epoch is set at the beginning of each training epoch (sampler.set_epoch(epoch)) to ensure different shuffling per epoch. The drop_last=True flag ensures all GPUs process the same number of batches.

NPOMTrainerDDP Class Structure

The NPOMTrainerDDP class encapsulates the entire training lifecycle:

NPOMTrainerDDP
├── __init__()           # Setup DDP, models, optimizers, data loaders, core pool
├── _build_corrector()   # Instantiate CorrectorNet with SyncBatchNorm + DDP wrapper
├── _borrow_poms()       # Get physics cores from pool or create fresh
├── run_physics_multistep()  # Core forward: inject ICs → multi-step with corrections
├── combine_loss_components() # Weighted sum of per-field losses
├── run_single_train_batch()  # One batch: forward → loss → backward → step
├── run_single_test_batch()   # One batch: forward → loss (no backwards)
├── train_epoch()        # Full training epoch loop
├── test_epoch()         # Full test epoch loop
├── log_step()           # Write per-batch loss stats to CSV
├── cleanup()            # Close log file, destroy process group

CorePool: Reusing Physics Core Instances

Physics core instantiation is expensive — it involves initializing grid geometries, computing Coriolis parameters, allocating large arrays, and running the model's setup routines. Rather than creating fresh cores for every batch, NeuralPOM uses a CorePool that pre-allocates cores and resets their state between batches:

class CorePool:
    def __init__(self, core_cls, device, pool_size, im, jm):
        # Pre-allocate pool_size core instances
        self.cores = [build_initialized_core(...) for _ in range(pool_size)]
        # Capture clean-state snapshots of each core
        self.snapshots = [capture_core_state(core) for core in self.cores]

    def reset_batch(self, batch_size):
        # Restore each core to its clean state via snapshot
        for core, snapshot in zip(self.cores[:batch_size], self.snapshots[:batch_size]):
            restore_core_state(core, snapshot)
        return self.cores[:batch_size]

The restore_core_state function copies values from the snapshot back to the core, replacing each tensor with a fresh detached clone. This is critical: without detaching, autograd history from a previous batch could leak into the next batch's computation graph.

CorePool is enabled via the --reuse_core_pool flag. When disabled, fresh cores are instantiated for each batch (slower but useful for debugging).

Training Loop (Per Batch)

The complete forward-backward pass for one batch:

def run_single_train_batch(self, ics, refs):
    # 1. Borrow physics cores from pool (or create fresh)
    poms = self._borrow_poms()

    # 2. Run multi-step physics with encoder/decoder corrections
    preds = self.run_physics_multistep(poms, ics, n_steps=refs["u"].shape[1])

    # 3. Compute per-field supervision losses
    losses = self._compute_component_losses(preds, refs)

    # 4. Weighted sum → total loss
    total_loss = self.combine_loss_components(losses)

    # 5. Backward pass (gradients flow through physics + networks)
    self.optimizer.zero_grad(set_to_none=True)
    total_loss.backward()

    # 6. Gradient clipping
    if self.args.grad_clip_norm > 0:
        torch.nn.utils.clip_grad_norm_(
            self.trainable_params, max_norm=self.args.grad_clip_norm
        )

    # 7. Optimizer step
    self.optimizer.step()

    # 8. Reduce loss stats across GPUs and return
    return self._reduce_losses(total_loss, losses)

Step 2: Multi-Step Physics with Corrections

The run_physics_multistep method is the heart of the hybrid pipeline. For each of n_steps correction cycles:

def run_physics_multistep(self, poms, ics, n_steps):
    # Inject initial conditions into all cores
    for index, pom in enumerate(poms):
        inject_core_state(pom, ics, sample_index=index)

    seq = {key: [] for key in PROGNOSTIC_FIELDS}

    for _ in range(n_steps):
        # (a) ENCODER: pack state → network → apply correction
        if self.encoder is not None:
            packed = pack_pom_batch(poms, self.normalizer)  # (B, 101, H, W)
            corrections = self.encoder(packed)               # (B, 101, H, W)
            for i, pom in enumerate(poms):
                apply_correction_to_pom(pom, corrections[i], self.normalizer,
                                       scale_factor=self.args.correction_scale)

        # (b) DECODER (pre-compute): pack state → network
        if self.decoder is not None:
            packed = pack_pom_batch(poms, self.normalizer)
            corrections_dec = self.decoder(packed)

        # (c) PHYSICS: advance all cores by 2 sub-steps (1200 s total)
        for pom in poms:
            pom.time0 = pom.time
            pom.iint = 0
            pom.iend = self.steps_per_correction  # = 2
            pom.run_time_loop()

        # (d) DECODER: apply correction after physics
        if self.decoder is not None:
            for i, pom in enumerate(poms):
                apply_correction_to_pom(pom, corrections_dec[i], self.normalizer,
                                       scale_factor=self.args.correction_scale)

        # (e) Record state after this correction cycle
        for key in seq:
            seq[key].append(torch.stack([getattr(pom, key) for pom in poms]))

    # Stack into (B, T, ...) tensors
    return {key: torch.stack(values, dim=1) for key, values in seq.items()}

Key details:

• Encoder is applied before physics, decoder after — this gives each network a distinct role.
• Decoder correction is pre-computed before the physics step (step b) but applied after (step d). This is a minor optimization: the packed tensor is the same regardless of when we run the decoder forward pass.
• Each correction cycle advances 1200 s (2 physics sub-steps of 600 s each), controlled by target_time_per_step=1200.0 and physics_step_seconds=600.0.
• Gradient flows through all steps: run_time_loop() is differentiable, so the backward pass propagates gradients through every physics sub-step and both CorrectorNets.

Step 5: Gradient Flow

The gradient computation is end-to-end differentiable. The chain of operations is:

$$\frac{\partial \mathcal{L}}{\partial \theta }=\frac{\partial \mathcal{L}}{\partial {\mathbf{pred}}_T}\cdot \frac{\partial {\mathbf{pred}}_T}{\partial {\mathbf{state}}_T}\cdot \frac{\partial {\mathbf{state}}_T}{\partial {\mathbf{state}}_{T-1}}\cdots \frac{\partial {\mathbf{state}}_1}{\partial \theta }$$

where $\theta$ represents the parameters of both encoder and decoder CorrectorNets, and each $\frac{\partial {\mathbf{state}}_t}{\partial {\mathbf{state}}_{t-1}}$ includes the gradient through the physics model's time integration.

Step 6: Gradient Clipping

Gradient clipping is critical for stable training because the gradient path through many physics steps can produce large accumulated gradients. The default grad_clip_norm clips the global norm of all trainable parameters:

torch.nn.utils.clip_grad_norm_(self.trainable_params, max_norm=grad_clip_norm)

Optimizer and Learning Rate Schedule

The optimizer is Adam with a cosine annealing learning rate schedule:

self.optimizer = optim.Adam(trainable_params, lr=args.lr)
self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
    self.optimizer, T_max=args.epochs, eta_min=1e-7
)

The learning rate follows:

$$\eta (e)={\eta }_{min}+\frac{1}{2}({\eta }_{max}-{\eta }_{min})(1+\cos \left(\frac{e}{E}\pi \right))$$

where $e$ is the current epoch, $E$ is the total number of epochs, ${\eta }_{max}$ is the initial learning rate, and ${\eta }_{min}={10}^{-7}$.

Correction Scale

The correction_scale hyperparameter controls the magnitude of corrections applied by the network:

apply_correction_to_pom(pom, correction, normalizer, scale_factor=args.correction_scale)

Default value: ${10}^{-3}$. This small scale factor has two purposes:

1. Numerical stability: Large corrections can destabilize the physics model during early training when the network is untrained and produces noisy outputs.

2. Gradient scaling: The scale factor effectively scales the gradient flowing into the CorrectorNet. A smaller scale means the network must learn to produce larger-magnitude outputs to achieve the same physical effect, which changes the effective learning dynamics.

The correction is multiplied element-wise: ${\mathbf{field}}_{\text{new}}={\mathbf{field}}_{\text{old}}+\alpha \cdot \mathbf{correction}$, where $\alpha$ is correction_scale.

Loss Weights

The total loss is a weighted sum of per-field MSE losses. The default weights are:

Field	Weight	Rationale
u	12.0	Primary prognostic field
ub	18.0	Barotropic velocity (large-scale flow)
v	10.0	Primary prognostic field
vb	12.0	Barotropic velocity (large-scale flow)
t	1.0	Temperature (already mean-subtracted)
s	6.0	Salinity (important for density)
el	30.0	Sea surface height (small magnitude, needs upweighting)
ua	60.0	Depth-averaged velocity (small magnitude)
uab	60.0	Depth-averaged barotropic velocity
va	60.0	Depth-averaged velocity (small magnitude)
vab	60.0	Depth-averaged barotropic velocity

The weights reflect the typical magnitude of each field and its relative importance. Surface and depth-averaged fields have small absolute values, so they need higher weights to contribute meaningfully to the total loss. These weights can be overridden via --loss_weights.

Logging and Checkpointing

Per-Batch Logging

Only the global rank 0 process writes logs. Each batch produces a CSV line with:

Epoch,Step,Phase,TotalLoss,U_MSE,UB_MSE,V_MSE,VB_MSE,T_MSE,S_MSE,EL_MSE,UA_MSE,UAB_MSE,VA_MSE,VAB_MSE

All per-field losses are multiplied by loss_display_scale for numerical convenience (to avoid extremely small numbers in the log).

Loss Reduction Across GPUs

After each batch, the loss statistics are reduced across all GPUs using dist.all_reduce with SUM operation, then divided by world_size to obtain the global average:

dist.all_reduce(stacked_losses, op=dist.ReduceOp.SUM)
global_average = stacked_losses / self.world_size

Epoch-Level Print

Rank 0 prints the average total loss for each training and test epoch.

Training on 8x A100 80GB GPUs

NeuralPOM is designed to train on 8 NVIDIA A100 80GB GPUs with the following typical configuration:

Parameter	Typical Value
GPUs	8
Per-GPU batch size	1-2
Global batch size	8-16
Grid resolution (im × jm)	360 × 240
Physics core	POM (Princeton Ocean Model)
Training epochs	50-200
Initial learning rate	$1\times {10}^{-4}$
Correction scale	${10}^{-3}$
Gradient clip norm	1.0-10.0
Encoder	Enabled
Decoder	Enabled
Core pool	Enabled
Training time per epoch	~10-30 minutes

The memory bottleneck is the physics core, not the neural network. Each POM instance requires several gigabytes of GPU memory for its state arrays, and the backward pass requires storing intermediate states for all physics sub-steps.

Loss Log Format

The training log file is a comma-separated CSV with the following columns:

Column	Description
Epoch	Current epoch number (0-indexed)
Step	Batch index within epoch
Phase	Train or Test
TotalLoss	Combined weighted loss (including loss_display_scale)
U_MSE	MSE for baroclinic zonal velocity
UB_MSE	MSE for barotropic zonal velocity
V_MSE	MSE for baroclinic meridional velocity
VB_MSE	MSE for barotropic meridional velocity
T_MSE	MSE for temperature
S_MSE	MSE for salinity
EL_MSE	MSE for sea surface height
UA_MSE	MSE for depth-averaged zonal velocity
UAB_MSE	MSE for barotropic depth-averaged zonal velocity
VA_MSE	MSE for depth-averaged meridional velocity
VAB_MSE	MSE for barotropic depth-averaged meridional velocity

All per-field MSE values are multiplied by loss_display_scale (default: 1.0) for numerical readability. The total loss is:

$${\mathcal{L}}_{\text{total}}={\lambda }_d\cdot \sum _f{w}_f\cdot {\mathcal{L}}_f$$

where ${\lambda }_d$ is loss_display_scale, $w_f$ are the per-field weights, and ${\mathcal{L}}_f$ are the per-field losses.