Fix training accumulation, evaluation and checkpoint loading

Initialize and refresh evaluation caches safely; migrate checkpoint keys before filtering; correct accumulation tails and distributed validation control flow and metric aggregation. Use a trainable decoder in the pretraining example and align environment installation with the pinned pip requirements. Update installation documentation. Validated with focused CPU/Gloo and RTX 4070 Ti SUPER CUDA regression checks, including same-precision real-checkpoint inference comparisons. Full downstream training accuracy and GPU multi-card execution were not benchmarked. Only the twelve reviewed source, configuration and documentation files are updated; model weights, paired model configuration and original datasets are unchanged.

BioMod/README.md

@@ -93,7 +93,8 @@ conda env create -f environment.yml
 conda activate biomod
 ```
 
-Or install Python packages into an existing CUDA/PyTorch environment:
+The Conda file installs the GPU stack through the pinned PyTorch/PyG wheels
+in `requirements.txt`. Alternatively, use a fresh Python 3.9 environment:
 
 ```bash
 python -m pip install -r requirements.txt

BioMod/biomod.py

@@ -57,6 +57,7 @@
 import torch.distributed as dist
 import sys
 from pathlib import Path
+from contextlib import nullcontext
 
 SOURCE_DIR = Path(__file__).resolve().parent
 PROJECT_ROOT = SOURCE_DIR.parent
@@ -261,8 +262,82 @@ def restore_model_args_from_config(args, config):
         args.no_gene2vec = not bool(config["use_gene2vec"])
     return args
 
+def evaluate_pretraining(
+    model, data_loader, device, amp_dtype, pin_memory=False, on_batch=None,
+):
+    """Evaluate without DDP forward collectives and aggregate valid batches.
+
+    Each rank may skip non-finite batches independently because inference uses
+    the underlying module. Unexpected validation errors are reported collectively;
+    an empty validation result is an error rather than a zero loss.
+    """
+    evaluation_model = model
+    buffer_sync_model = None
+    while True:
+        if hasattr(evaluation_model, "_orig_mod"):
+            evaluation_model = evaluation_model._orig_mod
+        elif isinstance(evaluation_model, DDP):
+            buffer_sync_model = evaluation_model
+            evaluation_model = evaluation_model.module
+        elif isinstance(evaluation_model, nn.DataParallel):
+            evaluation_model = evaluation_model.module
+        else:
+            break
+    evaluation_model.eval()
+    if buffer_sync_model is not None and buffer_sync_model.broadcast_buffers:
+        # Preserve the first DDP evaluation forward's buffer synchronization.
+        # Reuse PyTorch's coalesced broadcast once, before rank-local skips.
+        buffer_sync_model._sync_buffers()
+    loss_sum = 0.0
+    valid_batches = 0
+    skipped_batches = 0
+    first_error = None
+    batch_index = -1
+    try:
+        with torch.no_grad():
+            for batch_index, data in enumerate(data_loader):
+                data = data.to(device, non_blocking=pin_memory)
+                if not bool(torch.isfinite(data).all().item()):
+                    skipped_batches += 1
+                    continue
+                with autocast(
+                    device_type=torch.device(device).type, dtype=amp_dtype,
+                    enabled=(amp_dtype != torch.float32),
+                ):
+                    loss, _, _ = evaluation_model(data)
+                if not bool(torch.isfinite(loss).item()):
+                    skipped_batches += 1
+                    continue
+                batch_loss = float(loss.item())
+                loss_sum += batch_loss
+                valid_batches += 1
+                if on_batch is not None:
+                    on_batch(batch_index, batch_loss)
+    except Exception as exc:
+        first_error = f"at or after batch {batch_index}: {type(exc).__name__}: {exc}"
+
+    statistics = torch.tensor(
+        [loss_sum, valid_batches, skipped_batches, int(first_error is not None)],
+        dtype=torch.float64, device=device,
+    )
+    if dist.is_available() and dist.is_initialized():
+        dist.all_reduce(statistics, op=dist.ReduceOp.SUM)
+    total_loss, total_valid, total_skipped, failed_ranks = statistics.tolist()
+    if failed_ranks:
+        detail = first_error or "a validation error occurred on another rank"
+        raise RuntimeError(f"Validation failed on {int(failed_ranks)} rank(s); {detail}")
+    if total_valid == 0:
+        raise RuntimeError(
+            f"Validation produced no finite batches ({int(total_skipped)} skipped); "
+            "no validation loss is available."
+        )
+    return total_loss / total_valid, int(total_valid), int(total_skipped)
+
+
 if __name__ == "__main__":
     args = parser.parse_args()
+    if args.grad_acc <= 0:
+        raise ValueError("grad_acc must be a positive integer")
     resume_model_config = None
     resolved_model_config_path = None
     if args.model_path != 'None':
@@ -686,7 +761,7 @@ def save_initialized_checkpoint_and_exit(model_obj):
         loader_kwargs["prefetch_factor"] = args.prefetch_factor
     train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=train_sampler, **loader_kwargs)
     val_loader = DataLoader(val_dataset, batch_size=VALID_BATCH_SIZE, sampler=val_sampler, **loader_kwargs)
-    updates_per_epoch = max(1, len(train_loader) // GRADIENT_ACCUMULATION)
+    updates_per_epoch = max(1, math.ceil(len(train_loader) / GRADIENT_ACCUMULATION))
 
     # Step 3.6: Wrap the model using DDP (distributed data parallelism)
     # DDP will copy the model to each GPU and synchronize the gradients on each backward
@@ -757,71 +832,56 @@ def any_rank_true(local_flag: bool) -> bool:
         train_loader.sampler.set_epoch(i)  # DDP requires setting a different random seed for each epoch
         model.train()  # Set to training mode
         dist.barrier()  # Synchronize all processes
-        running_loss = 0.0  # Accumulate the loss of all steps in the epoch
+        running_loss = 0.0  # Accumulate finite training microbatch losses.
+        valid_train_steps = 0
+        optimizer.zero_grad()
+        discard_window = False
 
         # Step 4.2: Start batch loop (gradient accumulation)
         # Each iteration processes a batch, shape=(batch_size, num_genes)=(3, 27874)
         for index, data in enumerate(train_loader):
             index += 1
-            data = data.to(device, non_blocking=pin_memory)  # Move data to GPU, shape=(3, 27874)
-
-            # Check if input data contains NaN or Inf
-            invalid_input_local = bool(torch.isnan(data).any().item() or torch.isinf(data).any().item())
+            if (index - 1) % GRADIENT_ACCUMULATION == 0:
+                window_size = min(GRADIENT_ACCUMULATION, len(train_loader) - index + 1)
+                discard_window = False
+            if discard_window:
+                continue
+            update_boundary = (
+                index % GRADIENT_ACCUMULATION == 0 or index == len(train_loader)
+            )
+            data = data.to(device, non_blocking=pin_memory)
+            invalid_input_local = not bool(torch.isfinite(data).all().item())
             if any_rank_true(invalid_input_local):
                 if is_master:
-                    print(f"Warning: Input data contains NaN or Inf at step {index}, skipping...")
+                    print(f"Warning: Non-finite input at step {index}; discarding its accumulation window.")
+                optimizer.zero_grad()
+                discard_window = True
                 continue
 
-            # Step 4.3: Gradient accumulation training
-            # GRADIENT_ACCUMULATION=60: Update parameters after accumulating 60 batches of gradients.
-            # Equivalent batch_size = batch_size * grad_acc = 3 * 60 = 180
-            if index % GRADIENT_ACCUMULATION != 0:
-                # Step 4.3.1: Gradient accumulation step (no parameter update is performed)
-                # model.no_sync(): DDP does not perform gradient synchronization to reduce communication overhead.
-                with model.no_sync():
-                    with autocast(device_type='cuda', dtype=ptdtype):
-                        # [Model forward propagation]
-                        # Input: data shape=(3, 27874)
-                        # Internal processing:
-                        # 1. Flat to Module: (3, 27874) -> (3, num_modules, module_avg_size)
-                        #   2. Module Embedding: -> (3, num_modules, 1024)
-                        # 3. Random Masking: mask out 75% of modules
-                        # 4. BioMod encoder: Process visible modules
-                        # 5. Reconstruction: Predict the masked module
-                        # Output: loss=MSE on masked module, preds=reconstructed value, mask=mask mark
-                        loss, preds, mask = model(data)
-                        
-                        # Check if loss is NaN or Inf
-                        invalid_loss_local = bool(torch.isnan(loss).item() or torch.isinf(loss).item())
-                    if any_rank_true(invalid_loss_local):
-                        if is_master:
-                            print(f"Warning: Loss is NaN or Inf at step {index}, skipping...")
-                        continue
-
-                    # [Backpropagation] Cumulative gradient (no parameter update)
-                    # loss / GRADIENT_ACCUMULATION: Gradient will automatically accumulate and eventually average
-                    scaler.scale(loss / GRADIENT_ACCUMULATION).backward()# cumulative gradient
-            else:
-                # Step 4.3.2: Parameter update step (performed every GRADIENT_ACCUMULATION step)
-                # This step will perform gradient synchronization and parameter update.
+            # The final partial window uses its actual size. All other windows
+            # retain the original loss/grad_acc normalization and DDP reduction.
+            synchronization = nullcontext() if update_boundary else model.no_sync()
+            with synchronization:
                 with autocast(device_type='cuda', dtype=ptdtype):
                     loss, preds, mask = model(data)
-                    
-                    # Check if loss is NaN or Inf
-                    invalid_loss_local = bool(torch.isnan(loss).item() or torch.isinf(loss).item())
+                scaled_loss = scaler.scale(loss / window_size)
+                invalid_loss_local = not bool(torch.isfinite(loss).item())
                 if any_rank_true(invalid_loss_local):
-                    if is_master:
-                        print(f"Warning: Loss is NaN or Inf at step {index}, skipping...")
+                    if update_boundary:
+                        # A synchronized DDP forward must finish its reducer
+                        # cycle even when this window will not update weights.
+                        scaled_loss.backward()
                     optimizer.zero_grad()
+                    if scaler.is_enabled():
+                        scaler.update(new_scale=scaler.get_scale() * scaler.get_backoff_factor())
+                    discard_window = True
+                    if is_master:
+                        print(f"Warning: Non-finite loss at step {index}; discarding its accumulation window.")
                     continue
+                scaled_loss.backward()
 
-                # [Backpropagation] Accumulated gradient (the last accumulation, parameters will be updated soon)
-                scaler.scale(loss / GRADIENT_ACCUMULATION).backward()
-
-                # Step 4.3.3: Gradient clipping and parameter update
-                scaler.unscale_(optimizer)  # Descale gradients so that they crop correctly
-
-                # Check if the gradient contains NaN or Inf and do gradient clipping
+            if update_boundary:
+                scaler.unscale_(optimizer)
                 grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), int(1e2))
                 grad_norm_value = float(grad_norm.item())
                 invalid_grad_local = (
@@ -836,18 +896,17 @@ def any_rank_true(local_flag: bool) -> bool:
                     optimizer.zero_grad()
                     continue
 
-                # Update parameters
-                scaler.step(optimizer)# Update model parameters
+                scaler.step(optimizer)
                 scaler.update()
-                optimizer.zero_grad()  # Clear the gradient and prepare for the next round of accumulation
+                optimizer.zero_grad()
                 if args.lr_step_within_epoch:
                     scheduler.step()
 
-                if is_master: 
+                if is_master:
                     current_time = time.time()
                     if current_time - last_save_time >= save_time_interval:
                         time_stamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
-                        save_scaler_ckpt(time_stamp, model, optimizer, scheduler, scaler, loss.item(), model_name, ckpt_dir) # 1/3 every save_time_interval hour (default 24 hours) timestamp {YYYY-MM-DD_HH-MM-SS} during training
+                        save_scaler_ckpt(time_stamp, model, optimizer, scheduler, scaler, loss.item(), model_name, ckpt_dir)
                         last_save_time = current_time
 
             if index % LOG_EVERY == 0:
@@ -868,86 +927,29 @@ def any_rank_true(local_flag: bool) -> bool:
                         save_scaler_ckpt(i, model, optimizer, scheduler, scaler, loss.item(), model_name, latest_ckpt_dir)
 
             running_loss += loss.item()
+            valid_train_steps += 1
 
             # ========== [Step 5: Verification during training] ==========
             # Verify every VALIDATE_STEP step (default 1000 steps)
             if index % VALIDATE_STEP == 0:
                 model.eval()  # Switch to evaluation mode
                 dist.barrier()
-                running_loss_val = 0.0  # Cumulative validation set loss
-                val_index_num = 0  # Verify batch count
-                try:
-                    with torch.no_grad():  # Do not calculate gradients
-                        for val_index, data in enumerate(val_loader):
-                            data = data.to(device, non_blocking=pin_memory)
-                            
-                            # Check input data
-                            if torch.isnan(data).any() or torch.isinf(data).any():
-                                continue
-
-                            try:
-                                # Step 5.1: Validation set forward propagation
-                                # Input shape: (valid_batch_size, num_genes) = (5, 27874)
-                                with autocast(device_type='cuda', dtype=ptdtype):
-                                    loss_val, preds, mask = model(data)
-                                    
-                                    # Check loss
-                                    if torch.isnan(loss_val) or torch.isinf(loss_val):
-                                        continue
-                                    
-                                running_loss_val += loss_val.item()
-                                val_index_num += 1
-
-                                # Step 5.2: Print verification progress regularly
-                                if val_index % LOG_EVERY == 0:
-                                    if is_master: 
-                                        current_time = time.time()
-                                        if current_time - last_save_time >= save_time_interval:
-                                            time_stamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
-                                            save_scaler_ckpt(time_stamp, model, optimizer, scheduler, scaler, loss.item(), model_name, ckpt_dir) # 2/3 timestamp {YYYY-MM-DD_HH-MM-SS} every save_time_interval hours (default 24 hours) during the verification process
-                                            last_save_time = current_time
-                            except Exception as e:
-                                if is_master:
-                                    print(f"Warning: Error in validation step {val_index}: {e}")
-                                continue
-                            finally:
-                                # Make sure to clear memory
-                                del data
-                                if 'preds' in locals():
-                                    del preds
-                                if 'mask' in locals():
-                                    del mask
-                    
-                    # Make sure all processes have completed the verification loop
-                    dist.barrier()
-
-                    # Step 5.3: Calculate average verification loss
-                    # Each process calculates its own average loss and then synchronizes it through all_reduce
-                    if val_index_num > 0:
-                        loss_val_all = running_loss_val / val_index_num
-                    else:
-                        loss_val_all = 0.0
-
-                    # Step 5.4: DDP multi-process synchronization
-                    # Sum the val_index_num of all processes to ensure consistency
-                    val_index_num_tensor = torch.tensor(val_index_num, dtype=torch.float32, device=device)
-                    dist.all_reduce(val_index_num_tensor, op=dist.ReduceOp.SUM)
-                    total_val_index_num = int(val_index_num_tensor.item())
-
-                    # If all processes have no valid data, skip reduce
-                    if total_val_index_num > 0:
-                        # get_reduced: Collect the losses of all processes and calculate the average
-                        loss_val_reduced = get_reduced(loss_val_all, local_rank, 0, world_size)
-                    else:
-                        loss_val_reduced = 0.0
-                        
-                except Exception as e:
-                    if is_master:
-                        print(f"Error in validation: {e}")
-                    loss_val_reduced = 0.0
-                finally:
-                    # Make sure all processes are in sync
-                    dist.barrier()
+                validation_save_state = {"last_save_time": last_save_time}
+
+                def on_validation_batch(batch_index, batch_loss):
+                    if is_master and batch_index % LOG_EVERY == 0:
+                        current_time = time.time()
+                        if current_time - validation_save_state["last_save_time"] >= save_time_interval:
+                            time_stamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+                            save_scaler_ckpt(time_stamp, model, optimizer, scheduler, scaler, loss.item(), model_name, ckpt_dir)
+                            validation_save_state["last_save_time"] = current_time
+
+                loss_val_reduced, _, skipped_val_batches = evaluate_pretraining(
+                    model, val_loader, device, ptdtype, pin_memory, on_validation_batch,
+                )
+                last_save_time = validation_save_state["last_save_time"]
+                if is_master and skipped_val_batches:
+                    print(f"Warning: Skipped {skipped_val_batches} non-finite validation batches across all ranks.")
 
                 # Step 5.5: Print validation results
                 if is_master:
@@ -959,7 +961,10 @@ def any_rank_true(local_flag: bool) -> bool:
 
         # ========== [Step 6: Epoch end processing] ==========
         # Step 6.1: Calculate and record epoch training loss
-        epoch_loss = running_loss / index  # The average loss of this process
+        optimizer.zero_grad()  # Never carry a discarded window into the next epoch.
+        if valid_train_steps == 0:
+            raise RuntimeError("The epoch produced no finite training microbatches.")
+        epoch_loss = running_loss / valid_train_steps
         epoch_loss = get_reduced(epoch_loss, local_rank, 0, world_size)  # Multi-process average
         if is_master:
             print(f'    ==  Epoch: {i} | Training Loss: {epoch_loss:.6f}')
@@ -972,77 +977,24 @@ def any_rank_true(local_flag: bool) -> bool:
         if i % VALIDATE_EVERY == 0:
             model.eval()  # Switch to evaluation mode
             dist.barrier()
-            running_loss = 0.0
-            val_index_num = 0
-            try:
-                with torch.no_grad():
-                    # Step 7.1: Traverse the validation set
-                    for index, data in enumerate(val_loader):
-                        index += 1
-                        data = data.to(device, non_blocking=pin_memory)
-
-                        # Check input data
-                        if torch.isnan(data).any() or torch.isinf(data).any():
-                            continue
-
-                        try:
-                            # Step 7.2: Verify forward propagation
-                            with autocast(device_type='cuda', dtype=ptdtype):
-                                loss, preds, mask = model(data)
-                                
-                                # Check loss
-                                if torch.isnan(loss) or torch.isinf(loss):
-                                    continue
-                            
-                            running_loss += loss.item()
-                            val_index_num += 1
-                            
-                            if index % LOG_EVERY == 0:
-                                if is_master:
-                                    print(f"- epoch {i} step {index} val_loss {loss.item()} -")
-                                if is_master:
-                                    current_time = time.time()
-                                    if current_time - last_save_time >= save_time_interval:
-                                        time_stamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
-                                        save_scaler_ckpt(time_stamp, model, optimizer, scheduler, scaler, loss.item(), model_name, ckpt_dir)# Every save_time_interval hour (default 24 hours) during the 3/3 epoch verification process timestamp {YYYY-MM-DD_HH-MM-SS}
-                                        last_save_time = current_time
-                        except Exception as e:
-                            if is_master:
-                                print(f"Warning: Error in validation step {index}: {e}")
-                            continue
-                        finally:
-                            # Clear memory and release verification batch data
-                            del data
-                            if 'preds' in locals():
-                                del preds
-                            if 'mask' in locals():
-                                del mask
-
-                    # Step 7.3: Ensure all processes complete verification
-                    dist.barrier()
-
-                    # Step 7.4: Calculate average verification loss
-                    if val_index_num > 0:
-                        val_loss = running_loss / val_index_num
-                    else:
-                        val_loss = 0.0
+            validation_save_state = {"last_save_time": last_save_time}
 
-                    # Step 7.5: DDP multi-process synchronization
-                    val_index_num_tensor = torch.tensor(val_index_num, dtype=torch.float32, device=device)
-                    dist.all_reduce(val_index_num_tensor, op=dist.ReduceOp.SUM)
-                    total_val_index_num = int(val_index_num_tensor.item())
+            def on_epoch_validation_batch(batch_index, batch_loss):
+                validation_step = batch_index + 1
+                if is_master and validation_step % LOG_EVERY == 0:
+                    print(f"- epoch {i} step {validation_step} val_loss {batch_loss} -")
+                    current_time = time.time()
+                    if current_time - validation_save_state["last_save_time"] >= save_time_interval:
+                        time_stamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+                        save_scaler_ckpt(time_stamp, model, optimizer, scheduler, scaler, batch_loss, model_name, ckpt_dir)
+                        validation_save_state["last_save_time"] = current_time
 
-                    if total_val_index_num > 0:
-                        val_loss = get_reduced(val_loss, local_rank, 0, world_size)
-                    else:
-                        val_loss = 0.0
-                        
-            except Exception as e:
-                    if is_master:
-                        print(f"Error in epoch validation: {e}")
-                    val_loss = 0.0
-            finally:
-                dist.barrier()
+            val_loss, _, skipped_val_batches = evaluate_pretraining(
+                model, val_loader, device, ptdtype, pin_memory, on_epoch_validation_batch,
+            )
+            last_save_time = validation_save_state["last_save_time"]
+            if is_master and skipped_val_batches:
+                print(f"Warning: Skipped {skipped_val_batches} non-finite validation batches across all ranks.")
 
             # Step 7.6: Print epoch validation results
             if is_master:

BioMod/cell_annotation/train.py

@@ -980,6 +980,8 @@ def train_one_epoch(
     use_ddp: bool,
     selector_sparsity_weight: float = 0.0,
 ) -> dict:
+    if grad_accum_steps < 1:
+        raise ValueError("grad_accum_steps must be positive")
     model.train()
     running_loss = 0.0
     sample_count = 0
@@ -989,13 +991,18 @@ def train_one_epoch(
 
     num_steps = len(data_loader)
     for step, (x, y) in enumerate(data_loader, start=1):
+        # The final partial window must use its actual microbatch count and
+        # synchronize gradients before every optimizer step, including in DDP.
+        window_start = ((step - 1) // grad_accum_steps) * grad_accum_steps
+        window_steps = min(grad_accum_steps, num_steps - window_start)
+        should_step = step % grad_accum_steps == 0 or step == num_steps
         x = x.to(device)
         y = y.to(device)
         x_model = preprocessor.reindex_tensor_to_model_space(x, reindex_state=reindex_state)
 
         sync_context = (
             model.no_sync()
-            if use_ddp and grad_accum_steps > 1 and step % grad_accum_steps != 0
+            if use_ddp and not should_step
             else contextlib.nullcontext()
         )
         with sync_context:
@@ -1004,9 +1011,9 @@ def train_one_epoch(
             if selector_sparsity_weight > 0:
                 model_raw = model.module if use_ddp else model
                 loss = loss + selector_sparsity_weight * model_raw.get_selector_sparsity_loss()
-            (loss / grad_accum_steps).backward()
+            (loss / window_steps).backward()
 
-        if step % grad_accum_steps == 0 or step == num_steps:
+        if should_step:
             if grad_clip > 0:
                 torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
             optimizer.step()

BioMod/command.txt

@@ -22,7 +22,7 @@ torchrun --standalone --nproc_per_node="${NPROC_PER_NODE:-1}" \
     --model_name biomod_pretraining \
     --fusion_method multiply \
     --loss_mode gene_mean \
-    --decoder_type simple \
+    --decoder_type cross_attn \
     --num_workers 12 \
     --prefetch_factor 4 \
     --dec_embed_dim 1024

BioMod/environment.yml

@@ -1,67 +1,11 @@
 name: biomod
 channels:
-  - pytorch
-  - nvidia
-  - pyg
   - conda-forge
-  - defaults
 dependencies:
   - python=3.9
-  - pip>=24.0
-
-  # GPU stack matching the current server environment.
-  # For a different CUDA driver/toolkit, adjust pytorch-cuda and the PyG notes
-  # in requirements.txt accordingly.
-  - pytorch=2.6.*
-  - torchvision=0.21.*
-  - torchaudio=2.6.*
-  - pytorch-cuda=12.4
-  - pyg=2.6.*
-
-  # Core scientific stack
-  - numpy=1.26.*
-  - scipy=1.13.*
-  - pandas>=2.2,<2.4
-  - scikit-learn=1.6.*
-  - h5py>=3.10
-  - anndata=0.10.*
-  - scanpy=1.10.*
-  - matplotlib=3.9.*
-  - seaborn=0.13.*
-  - tqdm>=4.66
-  - requests>=2.31
-  - pyyaml>=6.0
-  - networkx>=3.2
-  - statsmodels>=0.14
-  - umap-learn>=0.5
-  - numba
-  - openpyxl
-  - jupyterlab
-  - ipykernel
-
-  # Single-cell and graph-analysis helpers
-  - python-igraph>=0.11
-  - leidenalg>=0.10
-  - rdkit
-
+  - pip>=24.0,<26
+  # Install the pinned CUDA 12.4 / PyTorch 2.6 / PyG stack from wheels.
+  # Conda's PyG 2.6 builds require older PyTorch releases.
+  # Conda resolves this path relative to this environment file.
   - pip:
-      # BioMod model utilities
-      - einops==0.8.1
-      - local-attention==1.11.2
-      - hickle>=5.0.0
-      - dcor==0.6
-
-      # Optional task modules used by downstream scripts.
-      - scib>=1.1.5
-      - decoupler>=1.8.0
-      - gseapy>=1.1.0
-      - mygene>=3.2.2
-      - lifelines>=0.27.0
-
-      # Text-alignment and LLM-assisted annotation utilities.
-      - sentence-transformers>=2.7.0
-      - transformers>=4.40.0
-      - openai>=1.40.0
-
-      # Only needed if running legacy pickles on Python < 3.8.
-      - pickle5>=0.0.12; python_version < "3.8"
+      - -r requirements.txt

BioMod/gene_perturbation/components/inference.py

@@ -98,6 +98,110 @@ def evaluate(loader, model, uncertainty, device, real_num_genes = None):
     
     return results
 
+
+def evaluate_perturbation_sums(loader, model, uncertainty, device, num_valid_samples):
+    """Accumulate real validation cells by perturbation without retaining cell matrices.
+
+    The sequential distributed sampler pads only after this rank's real cells.
+    Every batch still runs forward so all DDP ranks execute the same number of
+    forwards; padding contributes neither values nor counts to the statistics.
+    """
+    if num_valid_samples < 0:
+        raise ValueError("num_valid_samples must be non-negative")
+    model.eval()
+    model.to(device)
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        # Match the rank-zero checkpoint while keeping training buffer sync off.
+        # All ranks enter this single coalesced broadcast before validation.
+        model._sync_buffers()
+    statistics = {}
+    accepted = 0
+    with torch.no_grad():
+        for batch in loader:
+            batch.to(device)
+            output = model(batch)
+            prediction = output[0] if uncertainty else output
+            keep = min(len(batch.pert), num_valid_samples - accepted)
+            if keep == 0:
+                continue
+            truth = batch.y
+            # Match evaluate(): each cell uses its own ordered DE-gene indices.
+            values = {
+                'pred': prediction[:keep].detach().cpu().numpy(),
+                'truth': truth[:keep].detach().cpu().numpy(),
+                'pred_de': torch.stack([
+                    prediction[i, batch.de_idx[i]] for i in range(keep)
+                ]).detach().cpu().numpy(),
+                'truth_de': torch.stack([
+                    truth[i, batch.de_idx[i]] for i in range(keep)
+                ]).detach().cpu().numpy(),
+            }
+            conditions = np.asarray(batch.pert[:keep])
+            for condition in np.unique(conditions):
+                selected = conditions == condition
+                sums = {
+                    key: value[selected].sum(axis=0, dtype=np.float64)
+                    for key, value in values.items()
+                }
+                count = int(selected.sum())
+                if condition not in statistics:
+                    statistics[condition] = {'count': count, **sums}
+                else:
+                    statistics[condition]['count'] += count
+                    for key, value in sums.items():
+                        statistics[condition][key] += value
+            accepted += keep
+    if accepted != num_valid_samples:
+        raise ValueError(f"Expected {num_valid_samples} real validation cells, got {accepted}")
+    return statistics
+
+
+def merge_perturbation_sums(rank_statistics):
+    """Restore global perturbation means, the sufficient input to compute_metrics."""
+    merged = {}
+    fields = ('pred', 'truth', 'pred_de', 'truth_de')
+    for statistics in rank_statistics:
+        for condition, values in statistics.items():
+            if values['count'] <= 0:
+                raise ValueError("Perturbation sample counts must be positive")
+            if condition not in merged:
+                merged[condition] = {
+                    'count': values['count'],
+                    **{key: np.array(values[key], dtype=np.float64, copy=True) for key in fields},
+                }
+            else:
+                merged[condition]['count'] += values['count']
+                for key in fields:
+                    if merged[condition][key].shape != values[key].shape:
+                        raise ValueError(f"Inconsistent {key} shape for perturbation {condition}")
+                    merged[condition][key] += values[key]
+    if not merged:
+        raise ValueError("Validation requires at least one real cell")
+    conditions = sorted(merged)
+    return {
+        'pert_cat': np.asarray(conditions),
+        **{key: np.stack([merged[condition][key] / merged[condition]['count']
+                         for condition in conditions]) for key in fields},
+    }
+
+
+def gather_perturbation_metrics(local_statistics):
+    """Gather condition-level sums once; rank zero computes the unchanged metrics.
+
+    Communication scales with the number of conditions, not the number of cells.
+    Non-master ranks return None after participating in the same collective.
+    """
+    distributed = torch.distributed.is_available() and torch.distributed.is_initialized()
+    if distributed and torch.distributed.get_world_size() > 1:
+        rank = torch.distributed.get_rank()
+        gathered = [None] * torch.distributed.get_world_size() if rank == 0 else None
+        torch.distributed.gather_object(local_statistics, gathered, dst=0)
+        if rank != 0:
+            return None
+    else:
+        gathered = [local_statistics]
+    return compute_metrics(merge_perturbation_sums(gathered))[0]
+
 #This evaluation is just as test. IT HAS INFORMATION LEAKAGE!
 def evaluate_with_nonzero_mask(loader, model, uncertainty, device):
     """

BioMod/gene_perturbation/components/model.py

@@ -53,6 +53,7 @@
         BioModEmbedding,
         BioMod,
         sanitize_biomod_config,
+        migrate_legacy_biomod_state_dict,
     )
 except ModuleNotFoundError:
     project_root = str(Path(__file__).resolve().parents[3])
@@ -62,6 +63,7 @@
         BioModEmbedding,
         BioMod,
         sanitize_biomod_config,
+        migrate_legacy_biomod_state_dict,
     )
 
 # ============================================================
@@ -256,7 +258,10 @@ def __init__(self, input_genes, scfm_genes,
         # ============================================================
         if scfm_ckpt != 'None':
             ckpt = torch.load(scfm_ckpt, map_location=device)
-            model_state_dict = ckpt.get('model_state_dict', ckpt)
+            # Normalize historical keys before checking membership or shapes.
+            model_state_dict = migrate_legacy_biomod_state_dict(
+                ckpt.get('model_state_dict', ckpt)
+            )
             current_state_dict = self.scfm.state_dict()
             filtered_state_dict = {}
             for k, v in model_state_dict.items():

BioMod/gene_perturbation/components/trainer.py

@@ -26,6 +26,7 @@
 """
 
 from copy import deepcopy
+from contextlib import nullcontext
 import argparse
 from time import time
 import sys, os
@@ -48,7 +49,8 @@
 # ============================================================
 from .model import GEARS_Model_BioMod
 from .inference import evaluate, compute_metrics, deeper_analysis, \
-                  non_dropout_analysis, compute_synergy_loss
+                  non_dropout_analysis, compute_synergy_loss, \
+                  evaluate_perturbation_sums, gather_perturbation_metrics
 from .utils import loss_fct, uncertainty_loss_fct, parse_any_pert, \
                   get_similarity_network, print_sys, GeneSimNetwork, \
                   create_cell_graph_dataset_for_prediction, get_mean_control, \
@@ -556,9 +558,14 @@ def train(self,
              ):
         # Simulate larger batches through gradient accumulation. This value affects optimization stability and equivalent learning rate.
         GRADIENT_ACCUMULATION = self.config['accumulation_steps']
+        if GRADIENT_ACCUMULATION < 1:
+            raise ValueError("accumulation_steps must be at least 1")
         test = (epochs == 1)
         train_loader = self.dataloader['train_loader']
         val_loader = self.dataloader['val_loader']
+        epoch_batches = min(len(train_loader), 2 * GRADIENT_ACCUMULATION) if test else len(train_loader)
+        if epoch_batches == 0:
+            raise ValueError("Training requires at least one batch per process")
             
         self.model = self.model
 
@@ -591,34 +598,31 @@ def train(self,
         for epoch in range(epochs):
             train_loader.sampler.set_epoch(epoch)
             self.model.train()
+            optimizer.zero_grad(set_to_none=True)
             dist.barrier()
             for step, batch in enumerate(train_loader):
-                if test and (step+1 > 2 * GRADIENT_ACCUMULATION):
+                if step >= epoch_batches:
                     break
                 batch.to(self.device)
                 y = batch.y
-                if (step+1) % GRADIENT_ACCUMULATION != 0:
-                    # no_sync: Avoid doing all-reduce for each micro-batch to save communication.
-                    with self.model.no_sync():
-                        pred = self.model(batch)
-                        loss = loss_fct(pred, y, batch.pert,
-                                    ctrl = self.ctrl_expression, 
-                                    dict_filter = self.dict_filter,
-                                    direction_lambda = self.config['direction_lambda'])
-                    
-                        loss.backward()
-                elif (step+1) % GRADIENT_ACCUMULATION == 0:
+                window_start = (step // GRADIENT_ACCUMULATION) * GRADIENT_ACCUMULATION
+                window_size = min(GRADIENT_ACCUMULATION, epoch_batches - window_start)
+                update_now = (step - window_start + 1) == window_size
+                # Normalize by the actual window length, including an epoch's tail.
+                # Keep the existing microbatch loss definition and clipping rule.
+                sync_context = nullcontext() if update_now else self.model.no_sync()
+                with sync_context:
                     pred = self.model(batch)
                     loss = loss_fct(pred, y, batch.pert,
                                 ctrl = self.ctrl_expression, 
                                 dict_filter = self.dict_filter,
                                 direction_lambda = self.config['direction_lambda'])
-                
-                    loss.backward()
+                    (loss / window_size).backward()
+                if update_now:
                     # Gradient clipping affects training stability; if it is too small, it will inhibit learning, and if it is too large, the gradient may explode.
                     nn.utils.clip_grad_value_(self.model.parameters(), clip_value=1.0)
                     optimizer.step()
-                    optimizer.zero_grad()
+                    optimizer.zero_grad(set_to_none=True)
                 
                 if (step + 1) % 50 == 0:
                     if self.is_master:
@@ -627,15 +631,13 @@ def train(self,
             dist.barrier()
             scheduler.step()
             if (epoch + 1) % valid_every == 0:
-                # Evaluate model performance on train and val set
-                with torch.no_grad():
-                    val_res = evaluate(val_loader, self.model, self.config['uncertainty'], self.device)
-                val_metrics, _ = compute_metrics(val_res)
-                del val_res
-                val_metrics_reduced = {}
-                for m in ['mse', 'mse_de', 'pearson', 'pearson_de']:
-                    val_metrics_reduced[m] = get_reduced(val_metrics[m], self.local_rank, 0, self.world_size)
-                val_metrics = val_metrics_reduced
+                # Exclude sampler padding before forming global perturbation means.
+                val_statistics = evaluate_perturbation_sums(
+                    val_loader, self.model, self.config['uncertainty'], self.device,
+                    val_loader.sampler.num_valid_samples,
+                )
+                val_metrics = gather_perturbation_metrics(val_statistics)
+                del val_statistics
                 # Print epoch performance
                 log = "Epoch {}: Validation Overall MSE: {:.4f}. Validation Top 20 DE MSE: {:.4f}."
 

BioMod/gene_perturbation/components/training_utils.py

@@ -583,9 +583,15 @@ def __init__(self, dataset, batch_size, world_size, rank=None, num_replicas=None
         self.batch_size = batch_size
         self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.batch_size / self.num_replicas)) * self.batch_size
         self.total_size = self.num_samples * self.num_replicas
+        # All synthetic padding follows the real samples in this sequential layout.
+        self.num_valid_samples = max(
+            0, min(self.num_samples, len(self.dataset) - self.rank * self.num_samples)
+        )
 
     def __iter__(self):
         indices = list(range(len(self.dataset)))
+        if not indices:
+            return iter(indices)
         # Add padding samples to make the data size divisible
         indices += [indices[-1]] * (self.total_size - len(indices))
         # Subsampling: Each process only takes the part it is responsible for

BioMod/models_biomod.py

@@ -427,9 +427,9 @@ def __init__(
         if self.pool_affine_min_batch_size < 1:
             raise ValueError("MODULE_POOL_AFFINE_MIN_BATCH_SIZE must be positive")
         # Derived model weights only; no cell-dependent features are cached.
-        # These buffers add no parameters or checkpoint entries.
-        self.register_buffer("_pool_affine_score_weight", torch.empty(0), persistent=False)
-        self.register_buffer("_pool_affine_score_bias", torch.empty(0), persistent=False)
+        # Keep this rank-local cache out of DDP buffer broadcasts and checkpoints.
+        self._pool_affine_score_weight = None
+        self._pool_affine_score_bias = None
         self._pool_affine_cache_signature = None
 
     @property
@@ -442,12 +442,20 @@ def _pool_affine_signature(self):
                    self.attn_score.weight, self.attn_score.bias)
         return tuple((id(t), t.data_ptr(), t._version, t.device, t.dtype) for t in tensors)
 
+    def _apply(self, fn, recurse=True):
+        result = super()._apply(fn, recurse=recurse)
+        # Derived tensors are rebuilt on the destination device and dtype.
+        self._pool_affine_score_weight = None
+        self._pool_affine_score_bias = None
+        self._pool_affine_cache_signature = None
+        return result
+
     @torch.no_grad()
     def prepare_pool_affine(self):
-        """Compose final affine and score AFTER checkpoint/device/dtype placement.
+        """Compose the final affine and score for the current parameters.
 
-        Call again after changing weights, dtype or device. Forward rejects a
-        stale cache rather than silently using outdated model-derived weights.
+        Inference prepares this cache lazily and refreshes it after parameter
+        updates or checkpoint/device/dtype changes. Call explicitly to warm up.
         """
         weight = self.proj[2].weight
         compose_dtype = torch.float64 if weight.dtype == torch.float64 else torch.float32
@@ -474,7 +482,7 @@ def _forward_pool_affine(
         BF16 reassociation is validated separately; it is not bitwise equivalent.
         """
         if self._pool_affine_cache_signature != self._pool_affine_signature():
-            raise RuntimeError("Pool-affine cache is stale; call prepare_pool_affine() after loading/placing weights")
+            self.prepare_pool_affine()
         batch_size = expr.shape[0]
         flat_expr = torch.index_select(expr, 1, self._bucketed_gene_idx)
         flat_gene_emb = self.gene_embed(self._bucketed_gene_idx)

BioMod/requirements.txt

@@ -2,8 +2,8 @@
 #
 # Recommended use:
 #   1. Prefer `conda env create -f environment.yml` for a clean GPU setup.
-#   2. Use this file when PyTorch/CUDA are already installed and you only need
-#      to fill in Python package dependencies.
+#   2. Or install this file into a fresh Python 3.9 environment with pip.
+#      This installs the pinned GPU stack as well as task dependencies.
 #
 # CUDA 12.4 / PyTorch 2.6 wheel sources used by the server environment:
 --extra-index-url https://download.pytorch.org/whl/cu124
@@ -27,6 +27,11 @@ statsmodels==0.14.5
 umap-learn==0.5.9.post2
 dcor==0.6
 
+# Dataset tables and tutorial notebooks (also used by environment.yml).
+openpyxl>=3.1,<4
+jupyterlab>=4,<5
+ipykernel>=6,<7
+
 # PyTorch and graph learning stack
 torch==2.6.0+cu124
 torchvision==0.21.0+cu124

README.md

@@ -45,7 +45,9 @@ conda activate biomod
 ```
 
 The environment specification uses Python 3.9, PyTorch 2.6, CUDA 12.4, and PyTorch
-Geometric 2.6. See [environment.yml](BioMod/environment.yml) and
+Geometric 2.6. Conda creates the Python environment, then installs the pinned
+GPU packages from the PyTorch/PyG wheel indexes through `requirements.txt`.
+See [environment.yml](BioMod/environment.yml) and
 [requirements.txt](BioMod/requirements.txt) for the complete dependency lists.
 Match the GPU packages to your system when adapting the environment.