Refactor MTL: DDP NCCL support

geneformer/mtl/__init__.py

@@ -1 +1,4 @@
-# ruff: noqa: F401
+# ruff: noqa: F401
+
+from . import eval_utils
+from . import utils

geneformer/mtl/collators.py

@@ -1,8 +1,8 @@
 # imports
 import torch
 import pickle
-from ..collator_for_classification import DataCollatorForGeneClassification
-from .. import TOKEN_DICTIONARY_FILE
+from geneformer.collator_for_classification import DataCollatorForGeneClassification
+from geneformer import TOKEN_DICTIONARY_FILE
 
 """Geneformer collator for multi-task cell classification."""
 

geneformer/mtl/data.py

@@ -1,126 +1,190 @@
 import os
-from .collators import DataCollatorForMultitaskCellClassification
-from .imports import *
-
-def validate_columns(dataset, required_columns, dataset_type):
-    """Ensures required columns are present in the dataset."""
-    missing_columns = [col for col in required_columns if col not in dataset.column_names]
-    if missing_columns:
-        raise KeyError(
-            f"Missing columns in {dataset_type} dataset: {missing_columns}. "
-            f"Available columns: {dataset.column_names}"
-        )
-
-
-def create_label_mappings(dataset, task_to_column):
-    """Creates label mappings for the dataset."""
-    task_label_mappings = {}
-    num_labels_list = []
-    for task, column in task_to_column.items():
-        unique_values = sorted(set(dataset[column]))
-        mapping = {label: idx for idx, label in enumerate(unique_values)}
-        task_label_mappings[task] = mapping
-        num_labels_list.append(len(unique_values))
-    return task_label_mappings, num_labels_list
-
-
-def save_label_mappings(mappings, path):
-    """Saves label mappings to a pickle file."""
-    with open(path, "wb") as f:
-        pickle.dump(mappings, f)
-
-
-def load_label_mappings(path):
-    """Loads label mappings from a pickle file."""
-    with open(path, "rb") as f:
-        return pickle.load(f)
-
-
-def transform_dataset(dataset, task_to_column, task_label_mappings, config, is_test):
-    """Transforms the dataset to the required format."""
-    transformed_dataset = []
-    cell_id_mapping = {}
-
-    for idx, record in enumerate(dataset):
-        transformed_record = {
-            "input_ids": torch.tensor(record["input_ids"], dtype=torch.long),
-            "cell_id": idx,  # Index-based cell ID
-        }
-
-        if not is_test:
-            label_dict = {
-                task: task_label_mappings[task][record[column]]
-                for task, column in task_to_column.items()
-            }
-        else:
-            label_dict = {task: -1 for task in config["task_names"]}
-
-        transformed_record["label"] = label_dict
-        transformed_dataset.append(transformed_record)
-        cell_id_mapping[idx] = record.get("unique_cell_id", idx)
-
-    return transformed_dataset, cell_id_mapping
+import pickle
+import torch
+from torch.utils.data import DataLoader, Dataset
+from datasets import load_from_disk
 
+from .collators import DataCollatorForMultitaskCellClassification
 
-def load_and_preprocess_data(dataset_path, config, is_test=False, dataset_type=""):
-    """Main function to load and preprocess data."""
-    try:
-        dataset = load_from_disk(dataset_path)
 
+class StreamingMultiTaskDataset(Dataset):
+    
+    def __init__(self, dataset_path, config, is_test=False, dataset_type=""):
+        """Initialize the streaming dataset."""
+        self.dataset = load_from_disk(dataset_path)
+        self.config = config
+        self.is_test = is_test
+        self.dataset_type = dataset_type
+        self.cell_id_mapping = {}
+        
         # Setup task and column mappings
-        task_names = [f"task{i+1}" for i in range(len(config["task_columns"]))]
-        task_to_column = dict(zip(task_names, config["task_columns"]))
-        config["task_names"] = task_names
-
-        label_mappings_path = os.path.join(
+        self.task_names = [f"task{i+1}" for i in range(len(config["task_columns"]))]
+        self.task_to_column = dict(zip(self.task_names, config["task_columns"]))
+        config["task_names"] = self.task_names
+        
+        # Check if unique_cell_id column exists in the dataset
+        self.has_unique_cell_ids = "unique_cell_id" in self.dataset.column_names
+        print(f"{'Found' if self.has_unique_cell_ids else 'No'} unique_cell_id column in {dataset_type} dataset")
+        
+        # Setup label mappings
+        self.label_mappings_path = os.path.join(
             config["results_dir"],
             f"task_label_mappings{'_val' if dataset_type == 'validation' else ''}.pkl"
         )
-
+        
         if not is_test:
-            validate_columns(dataset, task_to_column.values(), dataset_type)
-
-            # Create and save label mappings
-            task_label_mappings, num_labels_list = create_label_mappings(dataset, task_to_column)
-            save_label_mappings(task_label_mappings, label_mappings_path)
+            self._validate_columns()
+            self.task_label_mappings, self.num_labels_list = self._create_label_mappings()
+            self._save_label_mappings()
         else:
             # Load existing mappings for test data
-            task_label_mappings = load_label_mappings(label_mappings_path)
-            num_labels_list = [len(mapping) for mapping in task_label_mappings.values()]
-
-        # Transform dataset
-        transformed_dataset, cell_id_mapping = transform_dataset(
-            dataset, task_to_column, task_label_mappings, config, is_test
-        )
-
-        return transformed_dataset, cell_id_mapping, num_labels_list
-
-    except KeyError as e:
-        raise ValueError(f"Configuration error or dataset key missing: {e}")
-    except Exception as e:
-        raise RuntimeError(f"Error during data loading or preprocessing: {e}")
+            self.task_label_mappings = self._load_label_mappings()
+            self.num_labels_list = [len(mapping) for mapping in self.task_label_mappings.values()]
+    
+    def _validate_columns(self):
+        """Ensures required columns are present in the dataset."""
+        missing_columns = [col for col in self.task_to_column.values() 
+                          if col not in self.dataset.column_names]
+        if missing_columns:
+            raise KeyError(
+                f"Missing columns in {self.dataset_type} dataset: {missing_columns}. "
+                f"Available columns: {self.dataset.column_names}"
+            )
+    
+    def _create_label_mappings(self):
+        """Creates label mappings for the dataset."""
+        task_label_mappings = {}
+        num_labels_list = []
+        
+        for task, column in self.task_to_column.items():
+            unique_values = sorted(set(self.dataset[column]))
+            mapping = {label: idx for idx, label in enumerate(unique_values)}
+            task_label_mappings[task] = mapping
+            num_labels_list.append(len(unique_values))
+        
+        return task_label_mappings, num_labels_list
+    
+    def _save_label_mappings(self):
+        """Saves label mappings to a pickle file."""
+        with open(self.label_mappings_path, "wb") as f:
+            pickle.dump(self.task_label_mappings, f)
+    
+    def _load_label_mappings(self):
+        """Loads label mappings from a pickle file."""
+        with open(self.label_mappings_path, "rb") as f:
+            return pickle.load(f)
+    
+    def __len__(self):
+        return len(self.dataset)
+    
+    def __getitem__(self, idx):
+        record = self.dataset[idx]
+        
+        # Store cell ID mapping
+        if self.has_unique_cell_ids:
+            unique_cell_id = record["unique_cell_id"]
+            self.cell_id_mapping[idx] = unique_cell_id
+        else:
+            self.cell_id_mapping[idx] = f"cell_{idx}"
+        
+        # Create transformed record
+        transformed_record = {
+            "input_ids": torch.tensor(record["input_ids"], dtype=torch.long),
+            "cell_id": idx,
+        }
+        
+        # Add labels
+        if not self.is_test:
+            label_dict = {
+                task: self.task_label_mappings[task][record[column]]
+                for task, column in self.task_to_column.items()
+            }
+        else:
+            label_dict = {task: -1 for task in self.config["task_names"]}
+        
+        transformed_record["label"] = label_dict
+        
+        return transformed_record
 
 
-def preload_and_process_data(config):
-    """Preloads and preprocesses train and validation datasets."""
-    # Process train data and save mappings
-    train_data = load_and_preprocess_data(config["train_path"], config, dataset_type="train")
+def get_data_loader(dataset, batch_size, sampler=None, shuffle=True):
+    """Create a DataLoader with the given dataset and parameters."""
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        sampler=sampler,
+        shuffle=shuffle if sampler is None else False,
+        num_workers=0,
+        pin_memory=True,
+        collate_fn=DataCollatorForMultitaskCellClassification(),
+    )
 
-    # Process validation data and save mappings
-    val_data = load_and_preprocess_data(config["val_path"], config, dataset_type="validation")
 
-    # Validate that the mappings match
+def prepare_data_loaders(config, include_test=False):
+    """Prepare data loaders for training, validation, and optionally test."""
+    result = {}
+    
+    # Process train data
+    train_dataset = StreamingMultiTaskDataset(
+        config["train_path"], 
+        config, 
+        dataset_type="train"
+    )
+    result["train_loader"] = get_data_loader(train_dataset, config["batch_size"])
+    
+    # Store the cell ID mapping from the dataset
+    result["train_cell_mapping"] = {k: v for k, v in train_dataset.cell_id_mapping.items()}
+    print(f"Collected {len(result['train_cell_mapping'])} cell IDs from training dataset")
+    
+    result["num_labels_list"] = train_dataset.num_labels_list
+    
+    # Process validation data
+    val_dataset = StreamingMultiTaskDataset(
+        config["val_path"], 
+        config, 
+        dataset_type="validation"
+    )
+    result["val_loader"] = get_data_loader(val_dataset, config["batch_size"])
+    
+    # Store the complete cell ID mapping for validation
+    for idx in range(len(val_dataset)):
+        _ = val_dataset[idx]
+    
+    result["val_cell_mapping"] = {k: v for k, v in val_dataset.cell_id_mapping.items()}
+    print(f"Collected {len(result['val_cell_mapping'])} cell IDs from validation dataset")
+    
+    # Validate label mappings
     validate_label_mappings(config)
-
-    return (*train_data[:2], *val_data)  # Return train and val data along with mappings
+    
+    # Process test data if requested
+    if include_test and "test_path" in config:
+        test_dataset = StreamingMultiTaskDataset(
+            config["test_path"], 
+            config, 
+            is_test=True,
+            dataset_type="test"
+        )
+        result["test_loader"] = get_data_loader(test_dataset, config["batch_size"])
+        
+        for idx in range(len(test_dataset)):
+            _ = test_dataset[idx]
+        
+        result["test_cell_mapping"] = {k: v for k, v in test_dataset.cell_id_mapping.items()}
+        print(f"Collected {len(result['test_cell_mapping'])} cell IDs from test dataset")
+    
+    return result
 
 
 def validate_label_mappings(config):
     """Ensures train and validation label mappings are consistent."""
     train_mappings_path = os.path.join(config["results_dir"], "task_label_mappings.pkl")
     val_mappings_path = os.path.join(config["results_dir"], "task_label_mappings_val.pkl")
-    train_mappings = load_label_mappings(train_mappings_path)
-    val_mappings = load_label_mappings(val_mappings_path)
+    
+    with open(train_mappings_path, "rb") as f:
+        train_mappings = pickle.load(f)
+    
+    with open(val_mappings_path, "rb") as f:
+        val_mappings = pickle.load(f)
 
     for task_name in config["task_names"]:
         if train_mappings[task_name] != val_mappings[task_name]:
@@ -131,32 +195,43 @@ def validate_label_mappings(config):
             )
 
 
-def get_data_loader(preprocessed_dataset, batch_size):
-    """Creates a DataLoader with optimal settings."""
-    return DataLoader(
-        preprocessed_dataset,
-        batch_size=batch_size,
-        shuffle=True,
-        collate_fn=DataCollatorForMultitaskCellClassification(),
-        num_workers=os.cpu_count(),
-        pin_memory=True,
+# Legacy functions for backward compatibility
+def preload_and_process_data(config):
+    """Preloads and preprocesses train and validation datasets."""
+    data = prepare_data_loaders(config)
+    
+    return (
+        data["train_loader"].dataset,
+        data["train_cell_mapping"],
+        data["val_loader"].dataset,
+        data["val_cell_mapping"],
+        data["num_labels_list"]
     )
 
 
 def preload_data(config):
     """Preprocesses train and validation data for trials."""
-    train_loader = get_data_loader(*preload_and_process_data(config)[:2], config["batch_size"])
-    val_loader = get_data_loader(*preload_and_process_data(config)[2:4], config["batch_size"])
-    return train_loader, val_loader
+    data = prepare_data_loaders(config)
+    return data["train_loader"], data["val_loader"]
 
 
 def load_and_preprocess_test_data(config):
     """Loads and preprocesses test data."""
-    return load_and_preprocess_data(config["test_path"], config, is_test=True)
+    test_dataset = StreamingMultiTaskDataset(
+        config["test_path"], 
+        config, 
+        is_test=True,
+        dataset_type="test"
+    )
+    
+    return (
+        test_dataset,
+        test_dataset.cell_id_mapping,
+        test_dataset.num_labels_list
+    )
 
 
 def prepare_test_loader(config):
     """Prepares DataLoader for test data."""
-    test_dataset, cell_id_mapping, num_labels_list = load_and_preprocess_test_data(config)
-    test_loader = get_data_loader(test_dataset, config["batch_size"])
-    return test_loader, cell_id_mapping, num_labels_list
+    data = prepare_data_loaders(config, include_test=True)
+    return data["test_loader"], data["test_cell_mapping"], data["num_labels_list"]

geneformer/mtl/eval_utils.py

@@ -1,19 +1,16 @@
+import os
+import json
+import torch
 import pandas as pd
 
-from .imports import *  # noqa # isort:skip
-from .data import prepare_test_loader  # noqa # isort:skip
+from .data import prepare_test_loader
 from .model import GeneformerMultiTask
 
-
 def evaluate_test_dataset(model, device, test_loader, cell_id_mapping, config):
     task_pred_labels = {task_name: [] for task_name in config["task_names"]}
     task_pred_probs = {task_name: [] for task_name in config["task_names"]}
     cell_ids = []
 
-    # # Load task label mappings from pickle file
-    # with open(f"{config['results_dir']}/task_label_mappings.pkl", "rb") as f:
-    #     task_label_mappings = pickle.load(f)
-
     model.eval()
     with torch.no_grad():
         for batch in test_loader:
@@ -85,4 +82,4 @@ def load_and_evaluate_test_model(config):
     best_model.to(device)
 
     evaluate_test_dataset(best_model, device, test_loader, cell_id_mapping, config)
-    print("Evaluation completed.")
+    print("Evaluation completed.")

geneformer/mtl/imports.py

@@ -1,43 +0,0 @@
-import functools
-import gc
-import json
-import os
-import pickle
-import sys
-import warnings
-from enum import Enum
-from itertools import chain
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-import optuna
-import pandas as pd
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-from datasets import load_from_disk
-from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, roc_curve
-from sklearn.model_selection import train_test_split
-from sklearn.preprocessing import LabelEncoder
-from torch.utils.data import DataLoader
-from transformers import (
-    AdamW,
-    BatchEncoding,
-    BertConfig,
-    BertModel,
-    DataCollatorForTokenClassification,
-    SpecialTokensMixin,
-    get_cosine_schedule_with_warmup,
-    get_linear_schedule_with_warmup,
-    get_scheduler,
-)
-from transformers.utils import logging, to_py_obj
-
-from .collators import DataCollatorForMultitaskCellClassification
-
-# local modules
-from .data import get_data_loader, preload_and_process_data
-from .model import GeneformerMultiTask
-from .optuna_utils import create_optuna_study
-from .utils import save_model

geneformer/mtl/model.py

@@ -118,4 +118,4 @@ class GeneformerMultiTask(nn.Module):
                         f"Error during loss computation for task {task_id}: {e}"
                     )
 
-        return total_loss, logits, losses if labels is not None else logits
+        return total_loss, logits, losses if labels is not None else logits

geneformer/mtl/optuna_utils.py

@@ -1,27 +0,0 @@
-import optuna
-from optuna.integration import TensorBoardCallback
-
-
-def save_trial_callback(study, trial, trials_result_path):
-    with open(trials_result_path, "a") as f:
-        f.write(
-            f"Trial {trial.number}: Value (F1 Macro): {trial.value}, Params: {trial.params}\n"
-        )
-
-
-def create_optuna_study(objective, n_trials, trials_result_path, tensorboard_log_dir):
-    study = optuna.create_study(direction="maximize")
-
-    # init TensorBoard callback
-    tensorboard_callback = TensorBoardCallback(
-        dirname=tensorboard_log_dir, metric_name="F1 Macro"
-    )
-
-    # callback and TensorBoard callback
-    callbacks = [
-        lambda study, trial: save_trial_callback(study, trial, trials_result_path),
-        tensorboard_callback,
-    ]
-
-    study.optimize(objective, n_trials=n_trials, callbacks=callbacks)
-    return study

geneformer/mtl/train.py

@@ -1,380 +1,707 @@
 import os
-import random
-
-import numpy as np
 import pandas as pd
 import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
+import optuna
+import functools
+import time
 
-from .imports import *
 from .model import GeneformerMultiTask
-from .utils import calculate_task_specific_metrics, get_layer_freeze_range
-
-
-def set_seed(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-def initialize_wandb(config):
-    if config.get("use_wandb", False):
-        import wandb
-
-        wandb.init(project=config["wandb_project"], config=config)
-        print("Weights & Biases (wandb) initialized and will be used for logging.")
-    else:
-        print(
-            "Weights & Biases (wandb) is not enabled. Logging will use other methods."
-        )
-
-
-def create_model(config, num_labels_list, device):
-    model = GeneformerMultiTask(
-        config["pretrained_path"],
-        num_labels_list,
-        dropout_rate=config["dropout_rate"],
-        use_task_weights=config["use_task_weights"],
-        task_weights=config["task_weights"],
-        max_layers_to_freeze=config["max_layers_to_freeze"],
-        use_attention_pooling=config["use_attention_pooling"],
-    )
-    if config["use_data_parallel"]:
-        model = nn.DataParallel(model)
-    return model.to(device)
-
-
-def setup_optimizer_and_scheduler(model, config, total_steps):
-    optimizer = AdamW(
-        model.parameters(),
-        lr=config["learning_rate"],
-        weight_decay=config["weight_decay"],
-    )
-    warmup_steps = int(config["warmup_ratio"] * total_steps)
-
-    if config["lr_scheduler_type"] == "linear":
-        scheduler = get_linear_schedule_with_warmup(
-            optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps
-        )
-    elif config["lr_scheduler_type"] == "cosine":
-        scheduler = get_cosine_schedule_with_warmup(
-            optimizer,
-            num_warmup_steps=warmup_steps,
-            num_training_steps=total_steps,
-            num_cycles=0.5,
-        )
-
-    return optimizer, scheduler
-
-
-def train_epoch(
-    model, train_loader, optimizer, scheduler, device, config, writer, epoch
-):
-    model.train()
-    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{config['epochs']}")
-    for batch_idx, batch in enumerate(progress_bar):
-        optimizer.zero_grad()
-        input_ids = batch["input_ids"].to(device)
-        attention_mask = batch["attention_mask"].to(device)
-        labels = [
-            batch["labels"][task_name].to(device) for task_name in config["task_names"]
-        ]
-
-        loss, _, _ = model(input_ids, attention_mask, labels)
-        loss.backward()
-
-        if config["gradient_clipping"]:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), config["max_grad_norm"])
-
-        optimizer.step()
-        scheduler.step()
-
-        writer.add_scalar(
-            "Training Loss", loss.item(), epoch * len(train_loader) + batch_idx
-        )
-        if config.get("use_wandb", False):
-            import wandb
-
-            wandb.log({"Training Loss": loss.item()})
-
-        # Update progress bar
-        progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})
-
-    return loss.item()  # Return the last batch loss
-
-
-def validate_model(model, val_loader, device, config):
-    model.eval()
-    val_loss = 0.0
-    task_true_labels = {task_name: [] for task_name in config["task_names"]}
-    task_pred_labels = {task_name: [] for task_name in config["task_names"]}
-    task_pred_probs = {task_name: [] for task_name in config["task_names"]}
-
-    with torch.no_grad():
-        for batch in val_loader:
-            input_ids = batch["input_ids"].to(device)
-            attention_mask = batch["attention_mask"].to(device)
+from .utils import (
+    calculate_metrics,
+    get_layer_freeze_range,
+    set_seed,
+    initialize_wandb,
+    create_model,
+    setup_optimizer_and_scheduler,
+    save_model,
+    save_hyperparameters,
+    prepare_training_environment,
+    log_training_step,
+    log_validation_metrics,
+    save_validation_predictions,
+    setup_logging,
+    setup_distributed_environment,
+    train_distributed
+)
+
+
+class Trainer:
+    """Trainer class for multi-task learning"""
+    
+    def __init__(self, config):
+        self.config = config
+        self.device = None
+        self.model = None
+        self.optimizer = None
+        self.scheduler = None
+        self.writer = None
+        self.is_distributed = config.get("distributed_training", False)
+        self.local_rank = config.get("local_rank", 0)
+        self.is_main_process = not self.is_distributed or self.local_rank == 0
+        
+    def train_epoch(self, train_loader, epoch):
+        """Train the model for one epoch."""
+        epoch_start = time.time()
+        self.model.train()
+        
+        # For distributed training, we need to be aware of the global batch count
+        if self.is_distributed:
+            # Get world size for reporting
+            world_size = dist.get_world_size()
+            # Calculate total batches across all GPUs
+            total_batches_global = len(train_loader) * world_size if self.local_rank == 0 else len(train_loader)
+        else:
+            world_size = 1
+            total_batches_global = len(train_loader)
+        
+        progress_bar = None
+        if self.is_main_process:
+            # Use the global batch count for progress reporting in distributed mode
+            progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{self.config['epochs']}", 
+                               total=len(train_loader))
+            iterator = progress_bar
+            
+            # Report distributed training information
+            if self.is_distributed:
+                print(f"Distributed training: {world_size} GPUs, {len(train_loader)} batches per GPU, "
+                      f"{total_batches_global} total batches globally")
+        else:
+            iterator = train_loader
+            
+        batch_times = []
+        forward_times = []
+        backward_times = []
+        optimizer_times = []
+        
+        # Get gradient accumulation steps from config (default to 1 if not specified)
+        accumulation_steps = self.config.get("gradient_accumulation_steps", 1)
+        
+        # Zero gradients at the beginning
+        self.optimizer.zero_grad()
+        
+        # Track loss for the entire epoch
+        total_loss = 0.0
+        num_batches = 0
+        accumulated_loss = 0.0
+        
+        for batch_idx, batch in enumerate(iterator):
+            batch_start = time.time()
+            
+            input_ids = batch["input_ids"].to(self.device)
+            attention_mask = batch["attention_mask"].to(self.device)
             labels = [
-                batch["labels"][task_name].to(device)
-                for task_name in config["task_names"]
+                batch["labels"][task_name].to(self.device) for task_name in self.config["task_names"]
             ]
-            loss, logits, _ = model(input_ids, attention_mask, labels)
-            val_loss += loss.item()
-
-            for sample_idx in range(len(batch["input_ids"])):
-                for i, task_name in enumerate(config["task_names"]):
-                    true_label = batch["labels"][task_name][sample_idx].item()
-                    pred_label = torch.argmax(logits[i][sample_idx], dim=-1).item()
-                    pred_prob = (
-                        torch.softmax(logits[i][sample_idx], dim=-1).cpu().numpy()
-                    )
-                    task_true_labels[task_name].append(true_label)
-                    task_pred_labels[task_name].append(pred_label)
-                    task_pred_probs[task_name].append(pred_prob)
-
-    val_loss /= len(val_loader)
-    return val_loss, task_true_labels, task_pred_labels, task_pred_probs
-
-
-def log_metrics(task_metrics, val_loss, config, writer, epochs):
-    for task_name, metrics in task_metrics.items():
-        print(
-            f"{task_name} - Validation F1 Macro: {metrics['f1']:.4f}, Validation Accuracy: {metrics['accuracy']:.4f}"
-        )
-        if config.get("use_wandb", False):
-            import wandb
 
-            wandb.log(
-                {
-                    f"{task_name} Validation F1 Macro": metrics["f1"],
-                    f"{task_name} Validation Accuracy": metrics["accuracy"],
-                }
+            forward_start = time.time()
+            loss, _, _ = self.model(input_ids, attention_mask, labels)
+            
+            # Scale loss by accumulation steps for gradient accumulation
+            if accumulation_steps > 1:
+                loss = loss / accumulation_steps
+            
+            forward_end = time.time()
+            forward_times.append(forward_end - forward_start)
+            
+            # Track loss - store the unscaled loss for reporting
+            unscaled_loss = loss.item() * (1 if accumulation_steps == 1 else accumulation_steps)
+            total_loss += unscaled_loss
+            num_batches += 1
+            accumulated_loss += loss.item()  # For gradient accumulation tracking
+            
+            backward_start = time.time()
+            
+            # Use no_sync() for all but the last accumulation step to avoid unnecessary communication
+            if self.is_distributed and accumulation_steps > 1:
+                # If this is not the last accumulation step or the last batch
+                if (batch_idx + 1) % accumulation_steps != 0 and (batch_idx + 1) != len(train_loader):
+                    with self.model.no_sync():
+                        loss.backward()
+                else:
+                    loss.backward()
+            else:
+                # Non-distributed training or accumulation_steps=1
+                loss.backward()
+            
+            backward_end = time.time()
+            backward_times.append(backward_end - backward_start)
+
+            # Only update weights after accumulation_steps or at the end of the epoch
+            if (batch_idx + 1) % accumulation_steps == 0 or (batch_idx + 1) == len(train_loader):
+                if self.config["gradient_clipping"]:
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config["max_grad_norm"])
+
+                optimizer_start = time.time()
+                self.optimizer.step()
+                self.scheduler.step()
+                self.optimizer.zero_grad()
+                optimizer_end = time.time()
+                optimizer_times.append(optimizer_end - optimizer_start)
+                
+                # Log after optimizer step
+                if self.is_main_process:
+                    # Calculate running average loss
+                    avg_loss = total_loss / num_batches
+                    
+                    log_training_step(avg_loss, self.writer, self.config, epoch, len(train_loader), batch_idx)
+                    
+                    # Update progress bar with just the running average loss
+                    progress_bar.set_postfix({"loss": f"{avg_loss:.4f}"})
+                
+                accumulated_loss = 0.0
+            else:
+                optimizer_times.append(0)  # No optimizer step taken
+            
+            batch_end = time.time()
+            batch_times.append(batch_end - batch_start)
+        
+        epoch_end = time.time()
+        
+        # Calculate average loss for the epoch
+        epoch_avg_loss = total_loss / num_batches
+        
+        # If distributed, gather losses from all processes to compute global average
+        if self.is_distributed:
+            # Create a tensor to hold the loss
+            loss_tensor = torch.tensor([epoch_avg_loss], device=self.device)
+            # Gather losses from all processes
+            all_losses = [torch.zeros_like(loss_tensor) for _ in range(dist.get_world_size())]
+            dist.all_gather(all_losses, loss_tensor)
+            # Compute the global average loss across all processes
+            epoch_avg_loss = torch.mean(torch.stack(all_losses)).item()
+        
+        if self.is_main_process:
+            # douhble check if batch_size has already been adjusted for world_size in the config
+            # This avoids double-counting the effective batch size
+            per_gpu_batch_size = self.config['batch_size']
+            total_effective_batch = per_gpu_batch_size * accumulation_steps * world_size
+            
+            print(f"Epoch {epoch+1} timing:")
+            print(f"  Total epoch time: {epoch_end - epoch_start:.2f}s")
+            print(f"  Average batch time: {sum(batch_times)/len(batch_times):.4f}s")
+            print(f"  Average forward time: {sum(forward_times)/len(forward_times):.4f}s")
+            print(f"  Average backward time: {sum(backward_times)/len(backward_times):.4f}s")
+            print(f"  Average optimizer time: {sum([t for t in optimizer_times if t > 0])/max(1, len([t for t in optimizer_times if t > 0])):.4f}s")
+            print(f"  Gradient accumulation steps: {accumulation_steps}")
+            print(f"  Batch size per GPU: {per_gpu_batch_size}")
+            print(f"  Effective global batch size: {total_effective_batch}")
+            print(f"  Average training loss: {epoch_avg_loss:.4f}")
+            if self.is_distributed:
+                print(f"  Total batches processed across all GPUs: {total_batches_global}")
+                print(f"  Communication optimization: Using no_sync() for gradient accumulation")
+
+        return epoch_avg_loss  # Return the average loss for the epoch
+
+    def validate_model(self, val_loader):
+        val_start = time.time()
+        self.model.eval()
+        val_loss = 0.0
+        task_true_labels = {task_name: [] for task_name in self.config["task_names"]}
+        task_pred_labels = {task_name: [] for task_name in self.config["task_names"]}
+        task_pred_probs = {task_name: [] for task_name in self.config["task_names"]}
+        
+        val_cell_ids = {}
+        sample_counter = 0
+
+        batch_times = []
+        
+        # Print validation dataset size
+        if self.is_main_process:
+            print(f"Validation dataset size: {len(val_loader.dataset)} samples")
+            print(f"Number of validation batches: {len(val_loader)}")
+
+            if self.is_distributed:
+                world_size = dist.get_world_size()
+                print(f"Distributed validation: {world_size} GPUs")
+                if hasattr(val_loader, 'sampler') and hasattr(val_loader.sampler, 'num_samples'):
+                    samples_per_gpu = val_loader.sampler.num_samples
+                    print(f"Each GPU processes {samples_per_gpu} validation samples")
+                    print(f"Total validation samples processed: {samples_per_gpu * world_size}")
+        
+        with torch.no_grad():
+            for batch in val_loader:
+                batch_start = time.time()
+                input_ids = batch["input_ids"].to(self.device)
+                attention_mask = batch["attention_mask"].to(self.device)
+                labels = [
+                    batch["labels"][task_name].to(self.device)
+                    for task_name in self.config["task_names"]
+                ]
+                loss, logits, _ = self.model(input_ids, attention_mask, labels)
+                val_loss += loss.item()
+
+                if "cell_id" in batch:
+                    for i, cell_id in enumerate(batch["cell_id"]):
+                        # Store the actual index for later mapping to unique_cell_id
+                        val_cell_ids[sample_counter + i] = cell_id.item()
+                
+                for sample_idx in range(len(batch["input_ids"])):
+                    for i, task_name in enumerate(self.config["task_names"]):
+                        true_label = batch["labels"][task_name][sample_idx].item()
+                        pred_label = torch.argmax(logits[i][sample_idx], dim=-1).item()
+                        # Store the full probability distribution
+                        pred_prob = torch.softmax(logits[i][sample_idx], dim=-1).cpu().numpy().tolist()
+                        task_true_labels[task_name].append(true_label)
+                        task_pred_labels[task_name].append(pred_label)
+                        task_pred_probs[task_name].append(pred_prob)
+                
+                # Update current index for cell ID tracking
+                sample_counter += len(batch["input_ids"])
+                
+                batch_end = time.time()
+                batch_times.append(batch_end - batch_start)
+
+        # norm validation loss by the number of batches
+        val_loss /= len(val_loader)
+        
+        # distributed, gather results from all processes
+        if self.is_distributed:
+            # Create tensors to hold the local results
+            loss_tensor = torch.tensor([val_loss], device=self.device)
+            gathered_losses = [torch.zeros_like(loss_tensor) for _ in range(dist.get_world_size())]
+            dist.all_gather(gathered_losses, loss_tensor)
+            
+            # Compute average loss across all processes
+            val_loss = torch.mean(torch.cat(gathered_losses)).item()
+            
+            world_size = dist.get_world_size()
+            
+            if self.is_main_process:
+                print(f"Collected predictions from rank {self.local_rank}")
+                print(f"Number of samples processed by this rank: {sample_counter}")
+            
+        val_end = time.time()
+        
+        if self.is_main_process:
+            print(f"Validation timing:")
+            print(f"  Total validation time: {val_end - val_start:.2f}s")
+            print(f"  Average batch time: {sum(batch_times)/len(batch_times):.4f}s")
+            print(f"  Collected {len(val_cell_ids)} cell indices from validation data")
+            print(f"  Processed {sample_counter} total samples during validation")
+            
+            # Print number of samples per task
+            for task_name in self.config["task_names"]:
+                print(f"  Task {task_name}: {len(task_true_labels[task_name])} samples")
+        
+        return val_loss, task_true_labels, task_pred_labels, task_pred_probs, val_cell_ids
+
+    def train(self, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list):
+        """Train the model and return validation loss and trained model."""
+        if self.config.get("use_wandb", False) and self.is_main_process:
+            initialize_wandb(self.config)
+
+        # Create model
+        self.model = create_model(self.config, num_labels_list, self.device, self.is_distributed, self.local_rank)
+        
+        # Setup optimizer and scheduler
+        total_steps = len(train_loader) * self.config["epochs"]
+        self.optimizer, self.scheduler = setup_optimizer_and_scheduler(self.model, self.config, total_steps)
+
+        # Training loop
+        if self.is_main_process:
+            epoch_progress = tqdm(range(self.config["epochs"]), desc="Training Progress")
+        else:
+            epoch_progress = range(self.config["epochs"])
+        
+        best_val_loss = float('inf')
+        train_losses = []
+        
+        with setup_logging(self.config) as self.writer:
+            for epoch in epoch_progress:
+                if self.is_distributed:
+                    train_loader.sampler.set_epoch(epoch)
+                    
+                train_loss = self.train_epoch(train_loader, epoch)
+                train_losses.append(train_loss)
+                
+                # Run validation after each epoch if configured to do so
+                if self.config.get("validate_each_epoch", False):
+                    val_loss, _, _, _, _ = self.validate_model(val_loader)
+                    if val_loss < best_val_loss:
+                        best_val_loss = val_loss
+                    
+                    if self.is_main_process:
+                        epoch_progress.set_postfix({
+                            "train_loss": f"{train_loss:.4f}",
+                            "val_loss": f"{val_loss:.4f}",
+                            "best_val_loss": f"{best_val_loss:.4f}"
+                        })
+                else:
+                    if self.is_main_process:
+                        epoch_progress.set_postfix({
+                            "train_loss": f"{train_loss:.4f}"
+                        })
+
+            val_loss, task_true_labels, task_pred_labels, task_pred_probs, val_cell_ids = self.validate_model(val_loader)
+            task_metrics = calculate_metrics(labels=task_true_labels, preds=task_pred_labels, metric_type="task_specific")
+
+            if self.is_main_process:
+                log_validation_metrics(task_metrics, val_loss, self.config, self.writer, self.config["epochs"])
+
+                # Save validation predictions
+                save_validation_predictions(
+                    val_cell_ids, 
+                    task_true_labels, 
+                    task_pred_labels, 
+                    task_pred_probs, 
+                    {**self.config, "val_cell_mapping": val_cell_id_mapping}  # Include the mapping
+                )
+
+                if self.config.get("use_wandb", False):
+                    import wandb
+                    wandb.finish()
+
+                print(f"\nTraining Summary:")
+                print(f"  Final Training Loss: {train_losses[-1]:.4f}")
+                print(f"  Final Validation Loss: {val_loss:.4f}")
+                for task_name, metrics in task_metrics.items():
+                    print(f"  {task_name} - F1: {metrics['f1']:.4f}, Accuracy: {metrics['accuracy']:.4f}")
+            
+        return val_loss, self.model  # Return both the validation loss and the trained model
+
+    def setup(self, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list):
+        if self.is_distributed:
+            self.device = torch.device(f"cuda:{self.local_rank}")
+        else:
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        self.model = create_model(self.config, num_labels_list, self.device)
+        
+        # war model w DDP
+        if self.is_distributed:
+            self.model = DDP(self.model, device_ids=[self.local_rank])
+            
+            # communication hook to optimize gradient synchronization
+            from torch.distributed.algorithms.ddp_comm_hooks import default as comm_hooks
+            
+            # default hook which maintains full precision
+            self.model.register_comm_hook(
+                state=None,
+                hook=comm_hooks.allreduce_hook
             )
-
-    writer.add_scalar("Validation Loss", val_loss, epochs)
-    for task_name, metrics in task_metrics.items():
-        writer.add_scalar(f"{task_name} - Validation F1 Macro", metrics["f1"], epochs)
-        writer.add_scalar(
-            f"{task_name} - Validation Accuracy", metrics["accuracy"], epochs
+            
+            print(f"Rank {self.local_rank}: Registered communication hook for optimized gradient synchronization")
+
+            print(f"Rank {self.local_rank}: Using samplers created in distributed worker")
+            print(f"Rank {self.local_rank}: Training dataset has {len(train_loader.dataset)} samples")
+            if hasattr(train_loader, 'sampler') and hasattr(train_loader.sampler, 'num_samples'):
+                print(f"Rank {self.local_rank}: This GPU will process {train_loader.sampler.num_samples} training samples per epoch")
+            
+            if hasattr(val_loader, 'sampler') and hasattr(val_loader.sampler, 'num_samples'):
+                print(f"Rank {self.local_rank}: This GPU will process {val_loader.sampler.num_samples} validation samples")
+        
+        # Set up optimizer and scheduler
+        self.optimizer, self.scheduler = setup_optimizer_and_scheduler(
+            self.model, self.config, len(train_loader)
         )
 
-
-def save_validation_predictions(
-    val_cell_id_mapping,
-    task_true_labels,
-    task_pred_labels,
-    task_pred_probs,
-    config,
-    trial_number=None,
-):
-    if trial_number is not None:
-        trial_results_dir = os.path.join(config["results_dir"], f"trial_{trial_number}")
-        os.makedirs(trial_results_dir, exist_ok=True)
-        val_preds_file = os.path.join(trial_results_dir, "val_preds.csv")
-    else:
-        val_preds_file = os.path.join(config["results_dir"], "manual_run_val_preds.csv")
-
-    rows = []
-    for sample_idx in range(len(val_cell_id_mapping)):
-        row = {"Cell ID": val_cell_id_mapping[sample_idx]}
-        for task_name in config["task_names"]:
-            row[f"{task_name} True"] = task_true_labels[task_name][sample_idx]
-            row[f"{task_name} Pred"] = task_pred_labels[task_name][sample_idx]
-            row[f"{task_name} Probabilities"] = ",".join(
-                map(str, task_pred_probs[task_name][sample_idx])
+        if self.is_main_process and self.config.get("use_wandb", False):
+            initialize_wandb(self.config)
+        
+        return train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list
+
+
+def train_model(config, device, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list):
+    """Train a model with the given configuration and data."""
+    # Check if distributed training is enabled
+    if config.get("distributed_training", False):
+        # Check if we have multiple GPUs
+        if torch.cuda.device_count() > 1:
+            result = train_distributed(
+                Trainer, 
+                config, 
+                train_loader, 
+                val_loader, 
+                train_cell_id_mapping, 
+                val_cell_id_mapping, 
+                num_labels_list
             )
-        rows.append(row)
-
-    df = pd.DataFrame(rows)
-    df.to_csv(val_preds_file, index=False)
-    print(f"Validation predictions saved to {val_preds_file}")
+            if result is not None:
+                return result
+        else:
+            print("Distributed training requested but only one GPU found. Falling back to single GPU training.")
+            config["distributed_training"] = False
+    
+    # Non-distributed training
+    trainer = Trainer(config)
+    trainer.device = device
+    return trainer.train(train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list)
 
 
-def train_model(
-    config,
-    device,
+def objective(
+    trial,
     train_loader,
     val_loader,
     train_cell_id_mapping,
     val_cell_id_mapping,
     num_labels_list,
+    config,
+    device,
 ):
+    """Objective function for Optuna hyperparameter optimization."""
     set_seed(config["seed"])
     initialize_wandb(config)
 
-    model = create_model(config, num_labels_list, device)
-    total_steps = len(train_loader) * config["epochs"]
-    optimizer, scheduler = setup_optimizer_and_scheduler(model, config, total_steps)
-
-    log_dir = os.path.join(config["tensorboard_log_dir"], "manual_run")
-    writer = SummaryWriter(log_dir=log_dir)
-
-    epoch_progress = tqdm(range(config["epochs"]), desc="Training Progress")
-    for epoch in epoch_progress:
-        last_loss = train_epoch(
-            model, train_loader, optimizer, scheduler, device, config, writer, epoch
+    trial_config = config.copy()
+    
+    # Suggest hyperparameters for this trial
+    for param_name, param_config in config["hyperparameters"].items():
+        if param_name == "lr_scheduler_type":
+            trial_config[param_name] = trial.suggest_categorical(
+                param_name, param_config["choices"]
+            )
+        elif param_name == "task_weights" and config["use_task_weights"]:
+            weights = [
+                trial.suggest_float(
+                    f"task_weight_{i}",
+                    param_config["low"],
+                    param_config["high"],
+                )
+                for i in range(len(num_labels_list))
+            ]
+            weight_sum = sum(weights)
+            trial_config[param_name] = [w / weight_sum for w in weights]
+        elif "log" in param_config and param_config["log"]:
+            trial_config[param_name] = trial.suggest_float(
+                param_name, param_config["low"], param_config["high"], log=True
+            )
+        else:
+            trial_config[param_name] = trial.suggest_float(
+                param_name, param_config["low"], param_config["high"]
+            )
+    
+    # Set appropriate max layers to freeze based on pretrained model
+    if "max_layers_to_freeze" in trial_config:
+        freeze_range = get_layer_freeze_range(trial_config["pretrained_path"])
+        trial_config["max_layers_to_freeze"] = int(trial.suggest_int(
+            "max_layers_to_freeze", 
+            freeze_range["min"], 
+            freeze_range["max"]
+        ))
+
+    trial_config["run_name"] = f"trial_{trial.number}"
+    
+    # Handle distributed training for this trial
+    if trial_config.get("distributed_training", False) and torch.cuda.device_count() > 1:
+        manager = mp.Manager()
+        shared_dict = manager.dict()
+        
+        train_distributed(
+            Trainer,
+            trial_config, 
+            train_loader, 
+            val_loader, 
+            train_cell_id_mapping, 
+            val_cell_id_mapping, 
+            num_labels_list,
+            trial.number, 
+            shared_dict
+        )
+        
+        val_loss = shared_dict.get('val_loss', float('inf'))
+        task_metrics = shared_dict.get('task_metrics', {})
+        
+        trial.set_user_attr("model_state_dict", shared_dict.get('model_state_dict', {}))
+        trial.set_user_attr("task_weights", trial_config["task_weights"])
+        
+        if config.get("use_wandb", False):
+            import wandb
+            wandb.log({
+                "trial_number": trial.number,
+                "val_loss": val_loss,
+                **{f"{task_name}_f1": metrics["f1"] for task_name, metrics in task_metrics.items()},
+                **{f"{task_name}_accuracy": metrics["accuracy"] for task_name, metrics in task_metrics.items()},
+            })
+            wandb.finish()
+            
+        return val_loss
+    
+    with setup_logging(trial_config) as writer:
+        trainer = Trainer(trial_config)
+        trainer.device = device
+        trainer.writer = writer
+        
+        # Create model with trial hyperparameters
+        trainer.model = create_model(trial_config, num_labels_list, device)
+        total_steps = len(train_loader) * config["epochs"]
+        trainer.optimizer, trainer.scheduler = setup_optimizer_and_scheduler(trainer.model, trial_config, total_steps)
+
+        # Training loop
+        for epoch in range(config["epochs"]):
+            trainer.train_epoch(train_loader, epoch)
+
+        val_loss, task_true_labels, task_pred_labels, task_pred_probs, val_cell_ids = trainer.validate_model(val_loader)
+        task_metrics = calculate_metrics(labels=task_true_labels, preds=task_pred_labels, metric_type="task_specific")
+
+        # Log metrics
+        log_validation_metrics(task_metrics, val_loss, trial_config, writer, config["epochs"])
+
+        # Save validation predictions
+        save_validation_predictions(
+            val_cell_ids,
+            task_true_labels,
+            task_pred_labels,
+            task_pred_probs,
+            {**trial_config, "val_cell_mapping": val_cell_id_mapping},
+            trial.number,
         )
-        epoch_progress.set_postfix({"last_loss": f"{last_loss:.4f}"})
 
-    val_loss, task_true_labels, task_pred_labels, task_pred_probs = validate_model(
-        model, val_loader, device, config
-    )
-    task_metrics = calculate_task_specific_metrics(task_true_labels, task_pred_labels)
+        # Store model state dict and task weights in trial user attributes
+        trial.set_user_attr("model_state_dict", trainer.model.state_dict())
+        trial.set_user_attr("task_weights", trial_config["task_weights"])
 
-    log_metrics(task_metrics, val_loss, config, writer, config["epochs"])
-    writer.close()
+        # Report intermediate value to Optuna
+        trial.report(val_loss, config["epochs"])
+        if trial.should_prune():
+            raise optuna.TrialPruned()
 
-    save_validation_predictions(
-        val_cell_id_mapping, task_true_labels, task_pred_labels, task_pred_probs, config
-    )
+        if config.get("use_wandb", False):
+            import wandb
+            wandb.log(
+                {
+                    "trial_number": trial.number,
+                    "val_loss": val_loss,
+                    **{f"{task_name}_f1": metrics["f1"] for task_name, metrics in task_metrics.items()},
+                    **{f"{task_name}_accuracy": metrics["accuracy"] for task_name, metrics in task_metrics.items()},
+                    **{k: v for k, v in trial_config.items() if k in [
+                        "learning_rate", "warmup_ratio", "weight_decay", "dropout_rate",
+                        "lr_scheduler_type", "use_attention_pooling", "max_layers_to_freeze"
+                    ]},
+                }
+            )
+            wandb.finish()
 
-    if config.get("use_wandb", False):
-        import wandb
+    return val_loss
 
-        wandb.finish()
 
-    print(f"\nFinal Validation Loss: {val_loss:.4f}")
-    return val_loss, model  # Return both the validation loss and the trained model
+def run_manual_tuning(config):
+    """Run training with manually specified hyperparameters."""
+    (
+        device,
+        train_loader,
+        val_loader,
+        train_cell_id_mapping,
+        val_cell_id_mapping,
+        num_labels_list,
+    ) = prepare_training_environment(config)
 
+    print("\nManual hyperparameters being used:")
+    for key, value in config["manual_hyperparameters"].items():
+        print(f"{key}: {value}")
+    print() 
 
-def objective(
-    trial,
-    train_loader,
-    val_loader,
-    train_cell_id_mapping,
-    val_cell_id_mapping,
-    num_labels_list,
-    config,
-    device,
-):
-    set_seed(config["seed"])  # Set the seed before each trial
-    initialize_wandb(config)
+    # Update config with manual hyperparameters
+    for key, value in config["manual_hyperparameters"].items():
+        config[key] = value
 
-    # Hyperparameters
-    config["learning_rate"] = trial.suggest_float(
-        "learning_rate",
-        config["hyperparameters"]["learning_rate"]["low"],
-        config["hyperparameters"]["learning_rate"]["high"],
-        log=config["hyperparameters"]["learning_rate"]["log"],
-    )
-    config["warmup_ratio"] = trial.suggest_float(
-        "warmup_ratio",
-        config["hyperparameters"]["warmup_ratio"]["low"],
-        config["hyperparameters"]["warmup_ratio"]["high"],
-    )
-    config["weight_decay"] = trial.suggest_float(
-        "weight_decay",
-        config["hyperparameters"]["weight_decay"]["low"],
-        config["hyperparameters"]["weight_decay"]["high"],
-    )
-    config["dropout_rate"] = trial.suggest_float(
-        "dropout_rate",
-        config["hyperparameters"]["dropout_rate"]["low"],
-        config["hyperparameters"]["dropout_rate"]["high"],
-    )
-    config["lr_scheduler_type"] = trial.suggest_categorical(
-        "lr_scheduler_type", config["hyperparameters"]["lr_scheduler_type"]["choices"]
-    )
-    config["use_attention_pooling"] = trial.suggest_categorical(
-        "use_attention_pooling", [False]
+    # Train the model
+    val_loss, trained_model = train_model(
+        config,
+        device,
+        train_loader,
+        val_loader,
+        train_cell_id_mapping,
+        val_cell_id_mapping,
+        num_labels_list,
     )
 
-    if config["use_task_weights"]:
-        config["task_weights"] = [
-            trial.suggest_float(
-                f"task_weight_{i}",
-                config["hyperparameters"]["task_weights"]["low"],
-                config["hyperparameters"]["task_weights"]["high"],
-            )
-            for i in range(len(num_labels_list))
-        ]
-        weight_sum = sum(config["task_weights"])
-        config["task_weights"] = [
-            weight / weight_sum for weight in config["task_weights"]
-        ]
-    else:
-        config["task_weights"] = None
-
-    # Dynamic range for max_layers_to_freeze
-    freeze_range = get_layer_freeze_range(config["pretrained_path"])
-    config["max_layers_to_freeze"] = trial.suggest_int(
-        "max_layers_to_freeze",
-        freeze_range["min"],
-        freeze_range["max"]
-    )
+    print(f"\nValidation loss with manual hyperparameters: {val_loss}")
 
-    model = create_model(config, num_labels_list, device)
-    total_steps = len(train_loader) * config["epochs"]
-    optimizer, scheduler = setup_optimizer_and_scheduler(model, config, total_steps)
+    # Save the trained model - only if not using distributed training
+    # (distributed training saves the model in the worker)
+    if not config.get("distributed_training", False):
+        model_save_directory = os.path.join(
+            config["model_save_path"], "GeneformerMultiTask"
+        )
+        save_model(trained_model, model_save_directory)
+
+        # Save the hyperparameters
+        hyperparams_to_save = {
+            **config["manual_hyperparameters"],
+            "dropout_rate": config["dropout_rate"],
+            "use_task_weights": config["use_task_weights"],
+            "task_weights": config["task_weights"],
+            "max_layers_to_freeze": config["max_layers_to_freeze"],
+            "use_attention_pooling": config["use_attention_pooling"],
+        }
+        save_hyperparameters(model_save_directory, hyperparams_to_save)
 
-    log_dir = os.path.join(config["tensorboard_log_dir"], f"trial_{trial.number}")
-    writer = SummaryWriter(log_dir=log_dir)
+    return val_loss
 
-    for epoch in range(config["epochs"]):
-        train_epoch(
-            model, train_loader, optimizer, scheduler, device, config, writer, epoch
-        )
 
-    val_loss, task_true_labels, task_pred_labels, task_pred_probs = validate_model(
-        model, val_loader, device, config
+def run_optuna_study(config):
+    """Run hyperparameter optimization using Optuna."""
+    # Prepare training environment
+    (
+        device,
+        train_loader,
+        val_loader,
+        train_cell_id_mapping,
+        val_cell_id_mapping,
+        num_labels_list,
+    ) = prepare_training_environment(config)
+
+    # If manual hyperparameters are specified, use them instead of running Optuna
+    if config.get("use_manual_hyperparameters", False):
+        return run_manual_tuning(config)
+
+    # Create a partial function with fixed arguments for the objective
+    objective_with_config_and_data = functools.partial(
+        objective,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        train_cell_id_mapping=train_cell_id_mapping,
+        val_cell_id_mapping=val_cell_id_mapping,
+        num_labels_list=num_labels_list,
+        config=config,
+        device=device,
     )
-    task_metrics = calculate_task_specific_metrics(task_true_labels, task_pred_labels)
 
-    log_metrics(task_metrics, val_loss, config, writer, config["epochs"])
-    writer.close()
+    # Create and run the Optuna study
+    study = optuna.create_study(
+        direction="minimize",  # Minimize validation loss
+        study_name=config["study_name"],
+        # storage=config["storage"],
+        load_if_exists=True,
+    )
 
-    save_validation_predictions(
-        val_cell_id_mapping,
-        task_true_labels,
-        task_pred_labels,
-        task_pred_probs,
-        config,
-        trial.number,
+    study.optimize(objective_with_config_and_data, n_trials=config["n_trials"])
+
+    # After finding the best trial
+    best_params = study.best_trial.params
+    best_task_weights = study.best_trial.user_attrs["task_weights"]
+    print("Saving the best model and its hyperparameters...")
+
+    # Create a model with the best hyperparameters
+    best_model = GeneformerMultiTask(
+        config["pretrained_path"],
+        num_labels_list,
+        dropout_rate=best_params["dropout_rate"],
+        use_task_weights=config["use_task_weights"],
+        task_weights=best_task_weights,
+        max_layers_to_freeze=best_params.get("max_layers_to_freeze", 0),
+        use_attention_pooling=best_params.get("use_attention_pooling", False),
     )
 
-    trial.set_user_attr("model_state_dict", model.state_dict())
-    trial.set_user_attr("task_weights", config["task_weights"])
+    # Get the best model state dictionary
+    best_model_state_dict = study.best_trial.user_attrs["model_state_dict"]
 
-    trial.report(val_loss, config["epochs"])
+    best_model_state_dict = {
+        k.replace("module.", ""): v for k, v in best_model_state_dict.items()
+    }
 
-    if trial.should_prune():
-        raise optuna.TrialPruned()
+    best_model.load_state_dict(best_model_state_dict, strict=False)
 
-    if config.get("use_wandb", False):
-        import wandb
+    model_save_directory = os.path.join(
+        config["model_save_path"], "GeneformerMultiTask"
+    )
+    save_model(best_model, model_save_directory)
 
-        wandb.log(
-            {
-                "trial_number": trial.number,
-                "val_loss": val_loss,
-                **{
-                    f"{task_name}_f1": metrics["f1"]
-                    for task_name, metrics in task_metrics.items()
-                },
-                **{
-                    f"{task_name}_accuracy": metrics["accuracy"]
-                    for task_name, metrics in task_metrics.items()
-                },
-                **{
-                    k: v
-                    for k, v in config.items()
-                    if k
-                    in [
-                        "learning_rate",
-                        "warmup_ratio",
-                        "weight_decay",
-                        "dropout_rate",
-                        "lr_scheduler_type",
-                        "use_attention_pooling",
-                        "max_layers_to_freeze",
-                    ]
-                },
-            }
-        )
-        wandb.finish()
+    save_hyperparameters(model_save_directory, {**best_params, "task_weights": best_task_weights})
 
-    return val_loss
+    return study.best_trial.value  # Return the best validation loss

geneformer/mtl/train_utils.py

@@ -1,161 +0,0 @@
-import random
-
-from .data import get_data_loader, preload_and_process_data
-from .imports import *
-from .model import GeneformerMultiTask
-from .train import objective, train_model
-from .utils import save_model
-
-
-def set_seed(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-def run_manual_tuning(config):
-    # Set seed for reproducibility
-    set_seed(config["seed"])
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    (
-        train_dataset,
-        train_cell_id_mapping,
-        val_dataset,
-        val_cell_id_mapping,
-        num_labels_list,
-    ) = preload_and_process_data(config)
-    train_loader = get_data_loader(train_dataset, config["batch_size"])
-    val_loader = get_data_loader(val_dataset, config["batch_size"])
-
-    # Print the manual hyperparameters being used
-    print("\nManual hyperparameters being used:")
-    for key, value in config["manual_hyperparameters"].items():
-        print(f"{key}: {value}")
-    print()  # Add an empty line for better readability
-
-    # Use the manual hyperparameters
-    for key, value in config["manual_hyperparameters"].items():
-        config[key] = value
-
-    # Train the model
-    val_loss, trained_model = train_model(
-        config,
-        device,
-        train_loader,
-        val_loader,
-        train_cell_id_mapping,
-        val_cell_id_mapping,
-        num_labels_list,
-    )
-
-    print(f"\nValidation loss with manual hyperparameters: {val_loss}")
-
-    # Save the trained model
-    model_save_directory = os.path.join(
-        config["model_save_path"], "GeneformerMultiTask"
-    )
-    save_model(trained_model, model_save_directory)
-
-    # Save the hyperparameters
-    hyperparams_to_save = {
-        **config["manual_hyperparameters"],
-        "dropout_rate": config["dropout_rate"],
-        "use_task_weights": config["use_task_weights"],
-        "task_weights": config["task_weights"],
-        "max_layers_to_freeze": config["max_layers_to_freeze"],
-        "use_attention_pooling": config["use_attention_pooling"],
-    }
-    hyperparams_path = os.path.join(model_save_directory, "hyperparameters.json")
-    with open(hyperparams_path, "w") as f:
-        json.dump(hyperparams_to_save, f)
-    print(f"Manual hyperparameters saved to {hyperparams_path}")
-
-    return val_loss
-
-
-def run_optuna_study(config):
-    # Set seed for reproducibility
-    set_seed(config["seed"])
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    (
-        train_dataset,
-        train_cell_id_mapping,
-        val_dataset,
-        val_cell_id_mapping,
-        num_labels_list,
-    ) = preload_and_process_data(config)
-    train_loader = get_data_loader(train_dataset, config["batch_size"])
-    val_loader = get_data_loader(val_dataset, config["batch_size"])
-
-    if config["use_manual_hyperparameters"]:
-        train_model(
-            config,
-            device,
-            train_loader,
-            val_loader,
-            train_cell_id_mapping,
-            val_cell_id_mapping,
-            num_labels_list,
-        )
-    else:
-        objective_with_config_and_data = functools.partial(
-            objective,
-            train_loader=train_loader,
-            val_loader=val_loader,
-            train_cell_id_mapping=train_cell_id_mapping,
-            val_cell_id_mapping=val_cell_id_mapping,
-            num_labels_list=num_labels_list,
-            config=config,
-            device=device,
-        )
-
-        study = optuna.create_study(
-            direction="minimize",  # Minimize validation loss
-            study_name=config["study_name"],
-            # storage=config["storage"],
-            load_if_exists=True,
-        )
-
-        study.optimize(objective_with_config_and_data, n_trials=config["n_trials"])
-
-        # After finding the best trial
-        best_params = study.best_trial.params
-        best_task_weights = study.best_trial.user_attrs["task_weights"]
-        print("Saving the best model and its hyperparameters...")
-
-        # Saving model as before
-        best_model = GeneformerMultiTask(
-            config["pretrained_path"],
-            num_labels_list,
-            dropout_rate=best_params["dropout_rate"],
-            use_task_weights=config["use_task_weights"],
-            task_weights=best_task_weights,
-        )
-
-        # Get the best model state dictionary
-        best_model_state_dict = study.best_trial.user_attrs["model_state_dict"]
-
-        # Remove the "module." prefix from the state dictionary keys if present
-        best_model_state_dict = {
-            k.replace("module.", ""): v for k, v in best_model_state_dict.items()
-        }
-
-        # Load the modified state dictionary into the model, skipping unexpected keys
-        best_model.load_state_dict(best_model_state_dict, strict=False)
-
-        model_save_directory = os.path.join(
-            config["model_save_path"], "GeneformerMultiTask"
-        )
-        save_model(best_model, model_save_directory)
-
-        # Additionally, save the best hyperparameters and task weights
-        hyperparams_path = os.path.join(model_save_directory, "hyperparameters.json")
-
-        with open(hyperparams_path, "w") as f:
-            json.dump({**best_params, "task_weights": best_task_weights}, f)
-        print(f"Best hyperparameters and task weights saved to {hyperparams_path}")

geneformer/mtl/utils.py

@@ -1,129 +1,641 @@
+from typing import Dict, List, Optional, Union
+import json
 import os
-import shutil
-
+import pickle
+import random
+import torch
+import numpy as np
+import wandb
+import optuna
 from sklearn.metrics import accuracy_score, f1_score
 from sklearn.preprocessing import LabelEncoder
-from transformers import AutoConfig, BertConfig, BertModel
+from torch.utils.tensorboard import SummaryWriter
+from transformers import AutoConfig, BertConfig, BertModel, get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup
+from torch.optim import AdamW
+import pandas as pd
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.multiprocessing as mp
+from contextlib import contextmanager
 
-from .imports import *
 
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def initialize_wandb(config):
+    if config.get("use_wandb", False):
+        wandb.init(
+            project=config.get("wandb_project", "geneformer_multitask"),
+            name=config.get("run_name", "experiment"),
+            config=config,
+            reinit=True,
+        )
+
+
+def create_model(config, num_labels_list, device, is_distributed=False, local_rank=0):
+    """Create and initialize the model based on configuration."""
+    from .model import GeneformerMultiTask
+    
+    model = GeneformerMultiTask(
+        config["pretrained_path"],
+        num_labels_list,
+        dropout_rate=config.get("dropout_rate", 0.1),
+        use_task_weights=config.get("use_task_weights", False),
+        task_weights=config.get("task_weights", None),
+        max_layers_to_freeze=config.get("max_layers_to_freeze", 0),
+        use_attention_pooling=config.get("use_attention_pooling", False),
+    )
+    
+    # Move model to device
+    model.to(device)
+    
+    if is_distributed:
+        model = DDP(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)
+    
+    return model
 
-def save_model(model, model_save_directory):
-    if not os.path.exists(model_save_directory):
-        os.makedirs(model_save_directory)
-
-    # Get the state dict
-    if isinstance(model, nn.DataParallel):
-        model_state_dict = (
-            model.module.state_dict()
-        )  # Use model.module to access the underlying model
-    else:
-        model_state_dict = model.state_dict()
 
-    # Remove the "module." prefix from the keys if present
-    model_state_dict = {
-        k.replace("module.", ""): v for k, v in model_state_dict.items()
+def setup_optimizer_and_scheduler(model, config, total_steps):
+    """Set up optimizer and learning rate scheduler."""
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() 
+                      if not any(nd in n for nd in no_decay) and p.requires_grad],
+            "weight_decay": config["weight_decay"],
+        },
+        {
+            "params": [p for n, p in model.named_parameters() 
+                      if any(nd in n for nd in no_decay) and p.requires_grad],
+            "weight_decay": 0.0,
+        },
+    ]
+    
+    optimizer = AdamW(
+        optimizer_grouped_parameters, 
+        lr=config["learning_rate"],
+        eps=config.get("adam_epsilon", 1e-8)
+    )
+    
+    # Prepare scheduler
+    warmup_steps = int(total_steps * config["warmup_ratio"])
+    
+    scheduler_map = {
+        "linear": get_linear_schedule_with_warmup,
+        "cosine": get_cosine_schedule_with_warmup
     }
+    
+    scheduler_fn = scheduler_map.get(config["lr_scheduler_type"])
+    if not scheduler_fn:
+        raise ValueError(f"Unsupported scheduler type: {config['lr_scheduler_type']}")
+    
+    scheduler = scheduler_fn(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)
+    
+    return optimizer, scheduler
+
+
+def save_model(model, model_save_directory):
+    """Save model weights and configuration."""
+    os.makedirs(model_save_directory, exist_ok=True)
+
+    # Handle DDP model
+    if isinstance(model, DDP):
+        model_to_save = model.module
+    else:
+        model_to_save = model
+    
+    model_state_dict = model_to_save.state_dict()
 
     model_save_path = os.path.join(model_save_directory, "pytorch_model.bin")
     torch.save(model_state_dict, model_save_path)
 
     # Save the model configuration
-    if isinstance(model, nn.DataParallel):
-        model.module.config.to_json_file(
-            os.path.join(model_save_directory, "config.json")
-        )
-    else:
-        model.config.to_json_file(os.path.join(model_save_directory, "config.json"))
+    model_to_save.config.to_json_file(os.path.join(model_save_directory, "config.json"))
 
     print(f"Model and configuration saved to {model_save_directory}")
 
 
-def calculate_task_specific_metrics(task_true_labels, task_pred_labels):
-    task_metrics = {}
-    for task_name in task_true_labels.keys():
-        true_labels = task_true_labels[task_name]
-        pred_labels = task_pred_labels[task_name]
-        f1 = f1_score(true_labels, pred_labels, average="macro")
-        accuracy = accuracy_score(true_labels, pred_labels)
-        task_metrics[task_name] = {"f1": f1, "accuracy": accuracy}
-    return task_metrics
+def save_hyperparameters(model_save_directory, hyperparams):
+    """Save hyperparameters to a JSON file."""
+    hyperparams_path = os.path.join(model_save_directory, "hyperparameters.json")
+    with open(hyperparams_path, "w") as f:
+        json.dump(hyperparams, f)
+    print(f"Hyperparameters saved to {hyperparams_path}")
 
 
-def calculate_combined_f1(combined_labels, combined_preds):
-    # Initialize the LabelEncoder
-    le = LabelEncoder()
+def calculate_metrics(labels=None, preds=None, task_data=None, metric_type="task_specific", return_format="dict"):
+    if metric_type == "single":
+        # Calculate metrics for a single task
+        if labels is None or preds is None:
+            raise ValueError("Labels and predictions must be provided for single task metrics")
+        
+        task_name = None
+        if isinstance(labels, dict) and len(labels) == 1:
+            task_name = list(labels.keys())[0]
+            labels = labels[task_name]
+            preds = preds[task_name]
+            
+        f1 = f1_score(labels, preds, average="macro")
+        accuracy = accuracy_score(labels, preds)
+        
+        if return_format == "tuple":
+            return f1, accuracy
+        
+        result = {"f1": f1, "accuracy": accuracy}
+        if task_name:
+            return {task_name: result}
+        return result
+        
+    elif metric_type == "task_specific":
+        # Calculate metrics for multiple tasks
+        if task_data:
+            result = {}
+            for task_name, (task_labels, task_preds) in task_data.items():
+                f1 = f1_score(task_labels, task_preds, average="macro")
+                accuracy = accuracy_score(task_labels, task_preds)
+                result[task_name] = {"f1": f1, "accuracy": accuracy}
+            return result
+        elif isinstance(labels, dict) and isinstance(preds, dict):
+            result = {}
+            for task_name in labels:
+                if task_name in preds:
+                    f1 = f1_score(labels[task_name], preds[task_name], average="macro")
+                    accuracy = accuracy_score(labels[task_name], preds[task_name])
+                    result[task_name] = {"f1": f1, "accuracy": accuracy}
+            return result
+        else:
+            raise ValueError("For task_specific metrics, either task_data or labels and preds dictionaries must be provided")
+            
+    elif metric_type == "combined":
+        # Calculate combined metrics across all tasks
+        if labels is None or preds is None:
+            raise ValueError("Labels and predictions must be provided for combined metrics")
+            
+        # Handle label encoding for non-numeric labels
+        if not all(isinstance(x, (int, float)) for x in labels + preds):
+            le = LabelEncoder()
+            le.fit(labels + preds)
+            labels = le.transform(labels)
+            preds = le.transform(preds)
+            
+        f1 = f1_score(labels, preds, average="macro")
+        accuracy = accuracy_score(labels, preds)
+        
+        if return_format == "tuple":
+            return f1, accuracy
+        return {"f1": f1, "accuracy": accuracy}
+        
+    else:
+        raise ValueError(f"Unknown metric_type: {metric_type}")
 
-    # Fit and transform combined labels and predictions to numerical values
-    le.fit(combined_labels + combined_preds)
-    encoded_true_labels = le.transform(combined_labels)
-    encoded_pred_labels = le.transform(combined_preds)
 
-    # Print out the mapping for sanity check
-    print("\nLabel Encoder Mapping:")
-    for index, class_label in enumerate(le.classes_):
-        print(f"'{class_label}': {index}")
+def get_layer_freeze_range(pretrained_path):
+    if not pretrained_path:
+        return {"min": 0, "max": 0}
+    
+    config = AutoConfig.from_pretrained(pretrained_path)
+    total_layers = config.num_hidden_layers
+    return {"min": 0, "max": total_layers - 1}
 
-    # Calculate accuracy
-    accuracy = accuracy_score(encoded_true_labels, encoded_pred_labels)
 
-    # Calculate F1 Macro score
-    f1 = f1_score(encoded_true_labels, encoded_pred_labels, average="macro")
+def prepare_training_environment(config):
+    """
+    Prepare the training environment by setting seed and loading data.
+    
+    Returns:
+        tuple: (device, train_loader, val_loader, train_cell_id_mapping, 
+               val_cell_id_mapping, num_labels_list)
+    """
+    from .data import prepare_data_loaders
+    
+    # Set seed for reproducibility
+    set_seed(config["seed"])
 
-    return f1, accuracy
+    # Set up device - for non-distributed training
+    if not config.get("distributed_training", False):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        # For distributed training, device will be set per process
+        device = None
+    
+    # Load data using the streaming dataset
+    data = prepare_data_loaders(config)
+    
+    # For distributed training, we'll set up samplers later in the distributed worker
+    # Don't create DistributedSampler here as process group isn't initialized yet
+    
+    return (
+        device,
+        data["train_loader"],
+        data["val_loader"],
+        data["train_cell_mapping"],
+        data["val_cell_mapping"],
+        data["num_labels_list"],
+    )
 
 
-# def save_model_without_heads(original_model_save_directory):
-#     # Create a new directory for the model without heads
-#     new_model_save_directory = original_model_save_directory + "_No_Heads"
-#     if not os.path.exists(new_model_save_directory):
-#         os.makedirs(new_model_save_directory)
+# Optuna hyperparameter optimization utilities
+def save_trial_callback(study, trial, trials_result_path):
+    """
+    Callback to save Optuna trial results to a file.
+    
+    Args:
+        study: Optuna study object
+        trial: Current trial object
+        trials_result_path: Path to save trial results
+    """
+    with open(trials_result_path, "a") as f:
+        f.write(
+            f"Trial {trial.number}: Value (F1 Macro): {trial.value}, Params: {trial.params}\n"
+        )
 
-#     # Load the model state dictionary
-#     model_state_dict = torch.load(
-#         os.path.join(original_model_save_directory, "pytorch_model.bin")
-#     )
 
-#     # Initialize a new BERT model without the classification heads
-#     config = BertConfig.from_pretrained(
-#         os.path.join(original_model_save_directory, "config.json")
-#     )
-#     model_without_heads = BertModel(config)
+def create_optuna_study(objective, n_trials: int, trials_result_path: str, tensorboard_log_dir: str) -> optuna.Study:
+    """Create and run an Optuna study with TensorBoard logging."""
+    from optuna.integration import TensorBoardCallback
+    
+    study = optuna.create_study(direction="maximize")
+    study.optimize(
+        objective,
+        n_trials=n_trials,
+        callbacks=[
+            lambda study, trial: save_trial_callback(study, trial, trials_result_path),
+            TensorBoardCallback(dirname=tensorboard_log_dir, metric_name="F1 Macro")
+        ]
+    )
+    return study
 
-#     # Filter the state dict to exclude classification heads
-#     model_without_heads_state_dict = {
-#         k: v
-#         for k, v in model_state_dict.items()
-#         if not k.startswith("classification_heads")
-#     }
 
-#     # Load the filtered state dict into the model
-#     model_without_heads.load_state_dict(model_without_heads_state_dict, strict=False)
+@contextmanager
+def setup_logging(config):
+    run_name = config.get("run_name", "manual_run")
+    log_dir = os.path.join(config["tensorboard_log_dir"], run_name)
+    writer = SummaryWriter(log_dir=log_dir)
+    try:
+        yield writer
+    finally:
+        writer.close()
 
-#     # Save the model without heads
-#     model_save_path = os.path.join(new_model_save_directory, "pytorch_model.bin")
-#     torch.save(model_without_heads.state_dict(), model_save_path)
 
-#     # Copy the configuration file
-#     shutil.copy(
-#         os.path.join(original_model_save_directory, "config.json"),
-#         new_model_save_directory,
-#     )
+def log_training_step(loss, writer, config, epoch, steps_per_epoch, batch_idx):
+    """Log training step metrics to TensorBoard and optionally W&B."""
+    writer.add_scalar(
+        "Training Loss", loss, epoch * steps_per_epoch + batch_idx
+    )
+    if config.get("use_wandb", False):
+        import wandb
+        wandb.log({"Training Loss": loss})
 
-#     print(f"Model without classification heads saved to {new_model_save_directory}")
 
+def log_validation_metrics(task_metrics, val_loss, config, writer, epoch):
+    """Log validation metrics to console, TensorBoard, and optionally W&B."""
+    for task_name, metrics in task_metrics.items():
+        print(
+            f"{task_name} - Validation F1 Macro: {metrics['f1']:.4f}, Validation Accuracy: {metrics['accuracy']:.4f}"
+        )
+        if config.get("use_wandb", False):
+            import wandb
+            wandb.log(
+                {
+                    f"{task_name} Validation F1 Macro": metrics["f1"],
+                    f"{task_name} Validation Accuracy": metrics["accuracy"],
+                }
+            )
 
-def get_layer_freeze_range(pretrained_path):
+    writer.add_scalar("Validation Loss", val_loss, epoch)
+    for task_name, metrics in task_metrics.items():
+        writer.add_scalar(f"{task_name} - Validation F1 Macro", metrics["f1"], epoch)
+        writer.add_scalar(
+            f"{task_name} - Validation Accuracy", metrics["accuracy"], epoch
+        )
+
+
+def load_label_mappings(results_dir: str, task_names: List[str]) -> Dict[str, Dict]:
+    """Load or initialize task label mappings."""
+    label_mappings_path = os.path.join(results_dir, "task_label_mappings_val.pkl")
+    if os.path.exists(label_mappings_path):
+        with open(label_mappings_path, 'rb') as f:
+            return pickle.load(f)
+    return {task_name: {} for task_name in task_names}
+
+
+def create_prediction_row(sample_idx: int, val_cell_indices: Dict, task_true_labels: Dict, 
+                         task_pred_labels: Dict, task_pred_probs: Dict, task_names: List[str], 
+                         inverted_mappings: Dict, val_cell_mapping: Dict) -> Dict:
+    """Create a row for validation predictions."""
+    batch_cell_idx = val_cell_indices.get(sample_idx)
+    cell_id = val_cell_mapping.get(batch_cell_idx, f"unknown_cell_{sample_idx}") if batch_cell_idx is not None else f"unknown_cell_{sample_idx}"
+    
+    row = {"Cell ID": cell_id}
+    for task_name in task_names:
+        if task_name in task_true_labels and sample_idx < len(task_true_labels[task_name]):
+            true_idx = task_true_labels[task_name][sample_idx]
+            pred_idx = task_pred_labels[task_name][sample_idx]
+            true_label = inverted_mappings.get(task_name, {}).get(true_idx, f"Unknown-{true_idx}")
+            pred_label = inverted_mappings.get(task_name, {}).get(pred_idx, f"Unknown-{pred_idx}")
+            
+            row.update({
+                f"{task_name}_true_idx": true_idx,
+                f"{task_name}_pred_idx": pred_idx,
+                f"{task_name}_true_label": true_label,
+                f"{task_name}_pred_label": pred_label
+            })
+            
+            if task_name in task_pred_probs and sample_idx < len(task_pred_probs[task_name]):
+                probs = task_pred_probs[task_name][sample_idx]
+                if isinstance(probs, (list, np.ndarray)) or (hasattr(probs, '__iter__') and not isinstance(probs, str)):
+                    prob_list = list(probs) if not isinstance(probs, list) else probs
+                    row[f"{task_name}_all_probs"] = ",".join(map(str, prob_list))
+                    for class_idx, prob in enumerate(prob_list):
+                        class_label = inverted_mappings.get(task_name, {}).get(class_idx, f"Unknown-{class_idx}")
+                        row[f"{task_name}_prob_{class_label}"] = prob
+                else:
+                    row[f"{task_name}_all_probs"] = str(probs)
+    
+    return row
+
+
+def save_validation_predictions(
+    val_cell_indices,
+    task_true_labels,
+    task_pred_labels,
+    task_pred_probs,
+    config,
+    trial_number=None,
+):
+    """Save validation predictions to a CSV file with class labels and probabilities."""
+    os.makedirs(config["results_dir"], exist_ok=True)
+    
+    if trial_number is not None:
+        os.makedirs(os.path.join(config["results_dir"], f"trial_{trial_number}"), exist_ok=True)
+        val_preds_file = os.path.join(config["results_dir"], f"trial_{trial_number}/val_preds.csv")
+    else:
+        val_preds_file = os.path.join(config["results_dir"], "manual_run_val_preds.csv")
+
+    if not val_cell_indices or not task_true_labels:
+        pd.DataFrame().to_csv(val_preds_file, index=False)
+        return
+
+    try:
+        label_mappings = load_label_mappings(config["results_dir"], config["task_names"])
+        inverted_mappings = {task: {idx: label for label, idx in mapping.items()} for task, mapping in label_mappings.items()}
+        val_cell_mapping = config.get("val_cell_mapping", {})
+        
+        # Determine maximum number of samples
+        max_samples = max(
+            [len(val_cell_indices)] + 
+            [len(task_true_labels[task]) for task in task_true_labels]
+        )
+        
+        rows = [
+            create_prediction_row(
+                sample_idx, val_cell_indices, task_true_labels, task_pred_labels, 
+                task_pred_probs, config["task_names"], inverted_mappings, val_cell_mapping
+            )
+            for sample_idx in range(max_samples)
+        ]
+        
+        pd.DataFrame(rows).to_csv(val_preds_file, index=False)
+    except Exception as e:
+        pd.DataFrame([{"Error": str(e)}]).to_csv(val_preds_file, index=False)
+
+
+def setup_distributed_environment(rank, world_size, config):
     """
-    Dynamically determines the number of layers to freeze based on the model depth from its configuration.
+    Setup the distributed training environment.
+    
     Args:
-        pretrained_path (str): Path to the pretrained model directory or model identifier.
-    Returns:
-        dict: A dictionary with 'min' and 'max' keys indicating the range of layers to freeze.
+        rank (int): The rank of the current process
+        world_size (int): Total number of processes
+        config (dict): Configuration dictionary
     """
-    if pretrained_path:
-        config = AutoConfig.from_pretrained(pretrained_path)
-        total_layers = config.num_hidden_layers
-        return {"min": 0, "max": total_layers - 1}
-    else:
-        return {"min": 0, "max": 0}
+    os.environ['MASTER_ADDR'] = config.get('master_addr', 'localhost')
+    os.environ['MASTER_PORT'] = config.get('master_port', '12355')
+    
+    # Initialize the process group
+    dist.init_process_group(
+        backend='nccl',
+        init_method='env://',
+        world_size=world_size,
+        rank=rank
+    )
+    
+    # Set the device for this process
+    torch.cuda.set_device(rank)
+
+
+def train_distributed(trainer_class, config, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list, trial_number=None, shared_dict=None):
+    """Run distributed training across multiple GPUs with fallback to single GPU."""
+    world_size = torch.cuda.device_count()
+    
+    if world_size <= 1:
+        print("Distributed training requested but only one GPU found. Falling back to single GPU training.")
+        config["distributed_training"] = False
+        trainer = trainer_class(config)
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        trainer.device = device
+        train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list = trainer.setup(
+            train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list
+        )
+        val_loss, model = trainer.train(
+            train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list
+        )
+        model_save_directory = os.path.join(config["model_save_path"], "GeneformerMultiTask")
+        save_model(model, model_save_directory)
+        save_hyperparameters(model_save_directory, {
+            **get_config_value(config, "manual_hyperparameters", {}),
+            "dropout_rate": config["dropout_rate"],
+            "use_task_weights": config["use_task_weights"],
+            "task_weights": config["task_weights"],
+            "max_layers_to_freeze": config["max_layers_to_freeze"],
+            "use_attention_pooling": config["use_attention_pooling"],
+        })
+        
+        if shared_dict is not None:
+            shared_dict['val_loss'] = val_loss
+            task_true_labels, task_pred_labels, task_pred_probs = collect_validation_predictions(model, val_loader, device, config)
+            shared_dict['task_metrics'] = calculate_metrics(labels=task_true_labels, preds=task_pred_labels, metric_type="task_specific")
+            shared_dict['model_state_dict'] = {k: v.cpu() for k, v in model.state_dict().items()}
+        
+        return val_loss, model
+    
+    print(f"Using distributed training with {world_size} GPUs")
+    mp.spawn(
+        _distributed_worker,
+        args=(world_size, trainer_class, config, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list, trial_number, shared_dict),
+        nprocs=world_size,
+        join=True
+    )
+    
+    if trial_number is None and shared_dict is None:
+        model_save_directory = os.path.join(config["model_save_path"], "GeneformerMultiTask")
+        model_path = os.path.join(model_save_directory, "pytorch_model.bin")
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        model = create_model(config, num_labels_list, device)
+        model.load_state_dict(torch.load(model_path))
+        return 0.0, model
+    
+    return None
+
+
+def _distributed_worker(rank, world_size, trainer_class, config, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list, trial_number=None, shared_dict=None):
+    """Worker function for distributed training."""
+    setup_distributed_environment(rank, world_size, config)
+    config["local_rank"] = rank
+    
+    # Set up distributed samplers
+    from torch.utils.data import DistributedSampler
+    from .data import get_data_loader
+    
+    train_sampler = DistributedSampler(train_loader.dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False)
+    val_sampler = DistributedSampler(val_loader.dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False)
+    
+    train_loader = get_data_loader(train_loader.dataset, config["batch_size"], sampler=train_sampler, shuffle=False)
+    val_loader = get_data_loader(val_loader.dataset, config["batch_size"], sampler=val_sampler, shuffle=False)
+    
+    if rank == 0:
+        print(f"Rank {rank}: Training {len(train_sampler)} samples, Validation {len(val_sampler)} samples")
+        print(f"Total samples across {world_size} GPUs: Training {len(train_sampler) * world_size}, Validation {len(val_sampler) * world_size}")
+    
+    # Create and setup trainer
+    trainer = trainer_class(config)
+    train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list = trainer.setup(
+        train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list
+    )
+    
+    # Train the model
+    val_loss, model = trainer.train(
+        train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list
+    )
+    
+    # Save model only from the main process
+    if rank == 0:
+        model_save_directory = os.path.join(config["model_save_path"], "GeneformerMultiTask")
+        save_model(model, model_save_directory)
+        
+        save_hyperparameters(model_save_directory, {
+            **get_config_value(config, "manual_hyperparameters", {}),
+            "dropout_rate": config["dropout_rate"],
+            "use_task_weights": config["use_task_weights"],
+            "task_weights": config["task_weights"],
+            "max_layers_to_freeze": config["max_layers_to_freeze"],
+            "use_attention_pooling": config["use_attention_pooling"],
+        })
+        
+        # For Optuna trials, store results in shared dictionary
+        if shared_dict is not None:
+            shared_dict['val_loss'] = val_loss
+            
+            # Run validation on full dataset from rank 0 for consistent metrics
+            full_val_loader = get_data_loader(val_loader.dataset, config["batch_size"], sampler=None, shuffle=False)
+            
+            # Get validation predictions using our utility function
+            task_true_labels, task_pred_labels, task_pred_probs = collect_validation_predictions(
+                model, full_val_loader, trainer.device, config
+            )
+            
+            # Calculate metrics
+            task_metrics = calculate_metrics(labels=task_true_labels, preds=task_pred_labels, metric_type="task_specific")
+            shared_dict['task_metrics'] = task_metrics
+            
+            # Store model state dict
+            if isinstance(model, DDP):
+                model_state_dict = model.module.state_dict()
+            else:
+                model_state_dict = model.state_dict()
+            
+            shared_dict['model_state_dict'] = {k: v.cpu() for k, v in model_state_dict.items()}
+    
+    # Clean up distributed environment
+    dist.destroy_process_group()
+
+
+def save_model_without_heads(model_directory):
+    """
+    Save a version of the fine-tuned model without classification heads.
+    
+    Args:
+        model_directory (str): Path to the directory containing the fine-tuned model
+    """
+    import torch
+    from transformers import BertConfig, BertModel
+    
+    # Load the full model
+    model_path = os.path.join(model_directory, "pytorch_model.bin")
+    config_path = os.path.join(model_directory, "config.json")
+    
+    if not os.path.exists(model_path) or not os.path.exists(config_path):
+        raise FileNotFoundError(f"Model files not found in {model_directory}")
+    
+    # Load the configuration
+    config = BertConfig.from_json_file(config_path)
+    
+    # Load the model state dict
+    state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+    
+    # Create a new model without heads
+    base_model = BertModel(config)
+    
+    # Filter out the classification head parameters
+    base_model_state_dict = {}
+    for key, value in state_dict.items():
+        # Only keep parameters that belong to the base model (not classification heads)
+        if not key.startswith('classification_heads') and not key.startswith('attention_pool'):
+            base_model_state_dict[key] = value
+    
+    # Load the filtered state dict into the base model
+    base_model.load_state_dict(base_model_state_dict, strict=False)
+    
+    # Save the model without heads
+    output_dir = os.path.join(model_directory, "model_without_heads")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Save the model weights
+    torch.save(base_model.state_dict(), os.path.join(output_dir, "pytorch_model.bin"))
+    
+    # Save the configuration
+    base_model.config.to_json_file(os.path.join(output_dir, "config.json"))
+    
+    print(f"Model without classification heads saved to {output_dir}")
+    return output_dir
+
+
+def get_config_value(config: Dict, key: str, default=None):
+    
+    return config.get(key, default)
+
+
+def collect_validation_predictions(model, val_loader, device, config) -> tuple:
+    task_true_labels = {}
+    task_pred_labels = {}
+    task_pred_probs = {}
+    
+    with torch.no_grad():
+        for batch in val_loader:
+            input_ids = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
+            labels = [batch["labels"][task_name].to(device) for task_name in config["task_names"]]
+            _, logits, _ = model(input_ids, attention_mask, labels)
+            
+            for sample_idx in range(len(batch["input_ids"])):
+                for i, task_name in enumerate(config["task_names"]):
+                    if task_name not in task_true_labels:
+                        task_true_labels[task_name] = []
+                        task_pred_labels[task_name] = []
+                        task_pred_probs[task_name] = []
+                    
+                    true_label = batch["labels"][task_name][sample_idx].item()
+                    pred_label = torch.argmax(logits[i][sample_idx], dim=-1).item()
+                    pred_prob = torch.softmax(logits[i][sample_idx], dim=-1).cpu().numpy()
+                    
+                    task_true_labels[task_name].append(true_label)
+                    task_pred_labels[task_name].append(pred_label)
+                    task_pred_probs[task_name].append(pred_prob)
+    
+    return task_true_labels, task_pred_labels, task_pred_probs

geneformer/mtl_classifier.py

@@ -29,7 +29,8 @@ Geneformer multi-task cell classifier.
 import logging
 import os
 
-from .mtl import eval_utils, train_utils, utils
+from .mtl import eval_utils, utils
+from .mtl.train import run_manual_tuning, run_optuna_study
 
 logger = logging.getLogger(__name__)
 
@@ -49,7 +50,9 @@ class MTLClassifier:
         "max_layers_to_freeze": {None, dict},
         "epochs": {None, int},
         "tensorboard_log_dir": {None, str},
-        "use_data_parallel": {None, bool},
+        "distributed_training": {None, bool},
+        "master_addr": {None, str},
+        "master_port": {None, str},
         "use_attention_pooling": {None, bool},
         "use_task_weights": {None, bool},
         "hyperparameters": {None, dict},
@@ -61,6 +64,7 @@ class MTLClassifier:
         "max_grad_norm": {None, int, float},
         "seed": {None, int},
         "trials_result_path": {None, str},
+        "gradient_accumulation_steps": {None, int},
     }
 
     def __init__(
@@ -79,7 +83,9 @@ class MTLClassifier:
         max_layers_to_freeze=None,
         epochs=1,
         tensorboard_log_dir="/results/tblogdir",
-        use_data_parallel=False,
+        distributed_training=False,
+        master_addr="localhost",
+        master_port="12355",
         use_attention_pooling=True,
         use_task_weights=True,
         hyperparameters=None,  # Default is None
@@ -89,6 +95,7 @@ class MTLClassifier:
         wandb_project=None,
         gradient_clipping=False,
         max_grad_norm=None,
+        gradient_accumulation_steps=1,  # Add this line with default value 1
         seed=42,  # Default seed value
     ):
         """
@@ -117,8 +124,12 @@ class MTLClassifier:
             | Path to directory to save results
         tensorboard_log_dir : None, str
             | Path to directory for Tensorboard logging results
-        use_data_parallel : None, bool
-            | Whether to use data parallelization
+        distributed_training : None, bool
+            | Whether to use distributed data parallel training across multiple GPUs
+        master_addr : None, str
+            | Master address for distributed training (default: localhost)
+        master_port : None, str
+            | Master port for distributed training (default: 12355)
         use_attention_pooling : None, bool
             | Whether to use attention pooling
         use_task_weights : None, bool
@@ -150,6 +161,8 @@ class MTLClassifier:
             | Whether to use gradient clipping
         max_grad_norm : None, int, float
             | Maximum norm for gradient clipping
+        gradient_accumulation_steps : None, int
+            | Number of steps to accumulate gradients before performing a backward/update pass
         seed : None, int
             | Random seed
         """
@@ -165,6 +178,7 @@ class MTLClassifier:
         self.batch_size = batch_size
         self.n_trials = n_trials
         self.study_name = study_name
+        self.gradient_accumulation_steps = gradient_accumulation_steps
 
         if max_layers_to_freeze is None:
             # Dynamically determine the range of layers to freeze
@@ -175,7 +189,9 @@ class MTLClassifier:
 
         self.epochs = epochs
         self.tensorboard_log_dir = tensorboard_log_dir
-        self.use_data_parallel = use_data_parallel
+        self.distributed_training = distributed_training
+        self.master_addr = master_addr
+        self.master_port = master_port
         self.use_attention_pooling = use_attention_pooling
         self.use_task_weights = use_task_weights
         self.hyperparameters = (
@@ -293,7 +309,7 @@ class MTLClassifier:
         self.config["manual_hyperparameters"] = self.manual_hyperparameters
         self.config["use_manual_hyperparameters"] = True
 
-        train_utils.run_manual_tuning(self.config)
+        run_manual_tuning(self.config)
 
     def validate_additional_options(self, req_var_dict):
         missing_variable = False
@@ -330,7 +346,7 @@ class MTLClassifier:
         req_var_dict = dict(zip(required_variable_names, required_variables))
         self.validate_additional_options(req_var_dict)
 
-        train_utils.run_optuna_study(self.config)
+        run_optuna_study(self.config)
 
     def load_and_evaluate_test_model(
         self,