Rank max learning rate

olmoearth_pretrain/evals/linear_probe.py

@@ -13,6 +13,7 @@
 import torch.nn.functional as F
 from einops import rearrange
 from olmo_core.data.utils import get_rng
+from olmo_core.distributed.utils import get_rank, get_world_size, is_distributed
 from sklearn.metrics import accuracy_score
 from torch.utils.data import DataLoader, TensorDataset
 from tqdm import tqdm
@@ -28,6 +29,20 @@
 logger = getLogger(__name__)
 
 
+RANK_MAX_LRS = [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1]
+
+
+def get_rank_lr(rank: int, world_size: int) -> float | None:
+    """Get the learning rate for this rank from the standard eval sweep LRs.
+
+    Uses the same LRs as eval sweeps: [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1]
+    Returns None if rank >= len(RANK_MAX_LRS) (rank should not participate).
+    """
+    if rank >= len(RANK_MAX_LRS):
+        return None
+    return RANK_MAX_LRS[rank]
+
+
 class ProbeType(StrEnum):
     """Enumeration of probe types for linear probing."""
 
@@ -135,15 +150,43 @@ def train_and_eval_probe(
     select_final_test_miou_based_on_epoch_of_max_val_miou: bool = False,
     n_bootstrap: int = 0,
     bootstrap_seed: int = 42,
+    rank_max_lr: bool = False,
 ) -> EvalTaskResult:
     """Run a linear probe on the OlmoEarth Pretrain model.
 
+    When rank_max_lr=True, each rank uses a different LR from a log-spaced range
+    and the best result (max primary metric) is returned via all_reduce.
+
     Returns:
         Dictionary with keys:
             - val_score: EvalResult for validation
             - test_score: EvalResult for test, or None if no test set
             - bootstrap_stats: Bootstrap statistics dict (empty dict if n_bootstrap == 0)
     """
+    # Rank-max LR: each rank uses a different LR from the standard sweep, then take max
+    # Auto-disable if world_size doesn't match the number of sweep LRs
+    use_rank_max = rank_max_lr and is_distributed()
+    if use_rank_max:
+        rank = get_rank()
+        world_size = get_world_size()
+        if world_size != len(RANK_MAX_LRS):
+            logger.warning(
+                f"rank_max_lr disabled: world_size={world_size} != {len(RANK_MAX_LRS)} "
+                f"(number of sweep LRs). Need exactly {len(RANK_MAX_LRS)} ranks."
+            )
+            use_rank_max = False
+
+    if use_rank_max:
+        rank_lr = get_rank_lr(rank, world_size)
+        assert rank_lr is not None
+        effective_lr = rank_lr
+        logger.info(
+            f"Rank-max LR enabled: rank {rank}/{world_size} using lr={effective_lr:.6f} "
+            f"(sweep LRs: {RANK_MAX_LRS})"
+        )
+    else:
+        effective_lr = lr
+
     logger.info(f"Probe type {probe_type}")
     if train_embeddings.shape[-1] != val_embeddings.shape[-1]:
         raise ValueError("Embedding dims don't match.")
@@ -204,7 +247,7 @@ def train_and_eval_probe(
             task_type=config.task_type,
             probe=probe,
             data_loader=data_loader,
-            lr=lr,
+            lr=effective_lr,
             epochs=end_epoch,
             total_epochs=epochs,
             current_epoch=start_epoch,

olmoearth_pretrain/train/callbacks/evaluator_callback.py

@@ -13,7 +13,7 @@
 import numpy as np
 import torch
 from olmo_core.train.callbacks.callback import Callback, CallbackConfig
-from olmo_core.train.common import Duration
+from olmo_core.train.common import Duration, ReduceType
 from olmo_core.train.trainer import Trainer
 from torch.utils.data import DataLoader
 
@@ -80,6 +80,7 @@ class DownstreamTaskConfig:
     probe_lr: float | None = None
     probe_batch_size: int = 32
     linear_probe_eval_interval: int = 50  # calculate val results every N epochs
+    rank_max_lr: bool = False  # each rank uses different LR, take max across ranks
     # FT
     ft_lr: float | None = None
     ft_batch_size: int = 32
@@ -144,6 +145,7 @@ def __init__(
         self.input_layers = task.input_layers
         self.probe_lr = task.probe_lr
         self.probe_batch_size = task.probe_batch_size
+        self.rank_max_lr = task.rank_max_lr
         self.ft_lr = task.ft_lr
         self.ft_batch_size = task.ft_batch_size
         self.finetune_seed = task.finetune_seed
@@ -213,6 +215,7 @@ def __init__(
                     probe_type=self.probe_type,
                     lr=self.probe_lr,
                     select_final_test_miou_based_on_epoch_of_max_val_miou=self.select_final_test_miou_based_on_epoch_of_max_val_miou,
+                    rank_max_lr=self.rank_max_lr,
                 )
                 if self.eval_mode == EvalMode.LINEAR_PROBE
                 else None
@@ -501,11 +504,20 @@ def _log_eval_result_to_wandb(
 def _record_eval_result(
     trainer: Trainer, prefix: str, name: str, result: EvalResult
 ) -> None:
-    """Record an EvalResult to trainer metrics."""
+    """Record an EvalResult to trainer metrics.
+
+    Uses ReduceType.max so that when ranks run different LRs (rank-max LR),
+    the best result across ranks is reported automatically.
+    """
     for metric_name, metric_value in result.metrics.items():
-        trainer.record_metric(f"{prefix}/{name}/{metric_name}", metric_value)
-    # Also record primary metric for backward compatibility
-    trainer.record_metric(f"{prefix}/{name}", result.primary)
+        trainer.record_metric(
+            f"{prefix}/{name}/{metric_name}",
+            metric_value,
+            reduce_type=ReduceType.max,
+        )
+    trainer.record_metric(
+        f"{prefix}/{name}", result.primary, reduce_type=ReduceType.max
+    )
 
 
 @dataclass

scripts/rank_max_lr_comparison.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# Comparison experiments for rank_max_lr feature
+# Run MADOS eval every 500 steps for 8000 steps total
+# Requires 8 GPUs for rank_max_lr to be active
+
+# Baseline (without rank_max_lr)
+python scripts/official/base.py launch mados-baseline ai2/jupiter-cirrascale-2 \
+  --trainer.max_duration.value=8000 \
+  --trainer.max_duration.unit=steps \
+  --trainer.callbacks.downstream_evaluator.tasks_to_run='["mados"]' \
+  --trainer.callbacks.downstream_evaluator.tasks.mados.eval_interval.value=500 \
+  --trainer.callbacks.downstream_evaluator.tasks.mados.eval_interval.unit=steps \
+  --trainer.callbacks.downstream_evaluator.tasks.mados.rank_max_lr=False \
+  --launch.num_gpus=8 \
+  --launch.priority=low
+
+# With rank_max_lr (each GPU uses different LR, take max)
+python scripts/official/base.py launch mados-rank-max-lr ai2/jupiter-cirrascale-2 \
+  --trainer.max_duration.value=8000 \
+  --trainer.max_duration.unit=steps \
+  --trainer.callbacks.downstream_evaluator.tasks_to_run='["mados"]' \
+  --trainer.callbacks.downstream_evaluator.tasks.mados.eval_interval.value=500 \
+  --trainer.callbacks.downstream_evaluator.tasks.mados.eval_interval.unit=steps \
+  --trainer.callbacks.downstream_evaluator.tasks.mados.rank_max_lr=True \
+  --launch.num_gpus=8 \
+  --launch.priority=low