diff --git a/olmoearth_pretrain/evals/linear_probe.py b/olmoearth_pretrain/evals/linear_probe.py index 3e728c130..60cfe8aec 100644 --- a/olmoearth_pretrain/evals/linear_probe.py +++ b/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, diff --git a/olmoearth_pretrain/train/callbacks/evaluator_callback.py b/olmoearth_pretrain/train/callbacks/evaluator_callback.py index e2101c08c..03545be38 100644 --- a/olmoearth_pretrain/train/callbacks/evaluator_callback.py +++ b/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 diff --git a/scripts/rank_max_lr_comparison.sh b/scripts/rank_max_lr_comparison.sh new file mode 100755 index 000000000..1a29246d1 --- /dev/null +++ b/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