From 0814e2fc35095821f8af35b5564b9775e5d0f25a Mon Sep 17 00:00:00 2001 From: "robin.cole@earthdaily.com" Date: Wed, 1 Jul 2026 14:12:12 +0100 Subject: [PATCH 1/2] Remove constraint on max_sequence_length and update related logic --- .../olmoearth_pretrain_v1/nn/flexi_vit.py | 46 ++++++++++------ .../olmoearth_pretrain_v1/utils/constants.py | 4 +- tests/test_flexi_vit.py | 53 ++++++++++++++++++- 3 files changed, 84 insertions(+), 19 deletions(-) diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py index 5ea16a4..6c81ec0 100644 --- a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py +++ b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py @@ -2,6 +2,7 @@ import logging import math +import warnings from dataclasses import dataclass from enum import StrEnum from typing import Any, NamedTuple @@ -840,7 +841,7 @@ def __init__( self, embedding_size: int, supported_modalities: list[ModalitySpec], - max_sequence_length: int, + max_sequence_length: int | None = None, learnable_channel_embeddings: bool = True, random_channel_embeddings: bool = False, tokenization_config: TokenizationConfig | None = None, @@ -853,7 +854,8 @@ def __init__( embedding_size: Size of token embeddings supported_modalities: Which modalities from Modality this model instantiation supports - max_sequence_length: Maximum sequence length + max_sequence_length: Deprecated, has no effect. Temporal position + encodings are now computed on-the-fly. learnable_channel_embeddings: Whether to use learnable channel embeddings random_channel_embeddings: Initialize channel embeddings randomly (zeros if False) tokenization_config: Optional config for custom band groupings @@ -870,6 +872,13 @@ def __init__( f"position_encoding must be one of {PositionEncoding.values()}, " f"got {position_encoding}" ) + if max_sequence_length is not None: + warnings.warn( + "max_sequence_length is deprecated and has no effect. " + "Temporal position encodings are now computed on-the-fly.", + DeprecationWarning, + stacklevel=2, + ) self.embedding_size = embedding_size self.supported_modalities = supported_modalities self.supported_modality_names = [ @@ -878,22 +887,15 @@ def __init__( self.tokenization_config = tokenization_config or TokenizationConfig() self.position_encoding = position_encoding self.embedding_size = embedding_size - self.max_sequence_length = ( - max_sequence_length # This max sequence length is a time dim thing - ) # TODO: we need to be able to calculate the size of the param based on what types of embeddings it will get # we have 4 embeddings types (pos_in_time, pos_in_space, month, channel) so each get # 0.25 of the dimension self.embedding_dim_per_embedding_type = int(embedding_size * 0.25) - # Position encodings for time dimension initialized to 1D sinusoidal encodings - self.pos_embed = nn.Parameter( - get_1d_sincos_pos_encoding( - torch.arange(max_sequence_length), - self.embedding_dim_per_embedding_type, - ), - requires_grad=False, - ) + # Temporal position encodings are computed on-the-fly via + # get_1d_sincos_pos_encoding so that any number of timesteps is supported + # without a pre-allocated table. + self._register_load_state_dict_pre_hook(self._drop_pos_embed_hook) # Month encodings month_tab = get_month_encoding_table(self.embedding_dim_per_embedding_type) self.month_embed = nn.Embedding.from_pretrained(month_tab, freeze=True) @@ -933,6 +935,15 @@ def _init_weights(self, m: nn.Module) -> None: # TODO: fix the dtype here nn.init.constant_(m.bias, 0).to(torch.float32) + @staticmethod + def _drop_pos_embed_hook( + state_dict: dict, prefix: str, *args: object, **kwargs: object + ) -> None: + """Drop legacy pos_embed from old checkpoints so strict loading succeeds.""" + key = prefix + "pos_embed" + if key in state_dict: + del state_dict[key] + @staticmethod def calculate_gsd_ratio(input_res: float, patch_size: int) -> float: """Calculate the Ground Sample Distance ratio.""" @@ -1032,10 +1043,13 @@ def _apply_encodings_per_modality( # Slot-index temporal encoding (additive). Skipped when 3D RoPE # handles temporal position rotationally inside attention. if not PositionEncoding.is_3d_rope(self.position_encoding): - time_embed = repeat( - self.pos_embed[:t], f"t d -> {ein_string}", **ein_dict + # Time position encodings (computed on-the-fly for arbitrary t) + pos_embed = get_1d_sincos_pos_encoding( + torch.arange(t, device=device), + self.embedding_dim_per_embedding_type, ) - modality_embed[..., n : n * 2] += time_embed.to(device) + time_embed = repeat(pos_embed, f"t d -> {ein_string}", **ein_dict) + modality_embed[..., n : n * 2] += time_embed # Month encodings stay additive in all modes (calendar/seasonal # signal is orthogonal to slot-index). diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py index 9f903a7..c35c040 100644 --- a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py +++ b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py @@ -4,7 +4,7 @@ """ from dataclasses import dataclass -from enum import Enum +from enum import StrEnum # The highest resolution that we are working at. # Everything else is a factor (which is a power of 2) coarser than this resolution. @@ -81,7 +81,7 @@ def get_expected_image_size(self, modality_resolution_factor: int) -> int: return IMAGE_TILE_SIZE // (self.resolution_factor // modality_resolution_factor) -class TimeSpan(str, Enum): +class TimeSpan(StrEnum): """Enum to distinguish data that is valid for different time ranges.""" # Only one data point (not time series). diff --git a/tests/test_flexi_vit.py b/tests/test_flexi_vit.py index faf7d25..8d4f624 100644 --- a/tests/test_flexi_vit.py +++ b/tests/test_flexi_vit.py @@ -1,7 +1,58 @@ -from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import Encoder +import pytest +import torch + +from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.encodings import ( + get_1d_sincos_pos_encoding, +) +from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import ( + CompositeEncodings, + Encoder, +) from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import Modality +def test_composite_encodings_compute_temporal_positions_on_the_fly() -> None: + """CompositeEncodings supports sequences longer than deprecated config value.""" + with pytest.warns(DeprecationWarning): + encodings = CompositeEncodings( + embedding_size=16, + supported_modalities=[Modality.SENTINEL2_L2A], + max_sequence_length=3, + ) + assert "pos_embed" not in encodings.state_dict() + + tokens = torch.zeros((1, 2, 2, 5, Modality.SENTINEL2_L2A.num_band_sets, 16)) + timestamps = torch.tensor( + [[[1, 0, 2020], [2, 1, 2020], [3, 2, 2020], [4, 3, 2020], [5, 4, 2020]]] + ) + + output = encodings( + {Modality.SENTINEL2_L2A.name: tokens}, + timestamps=timestamps, + patch_size=4, + ) + + encoded = output[Modality.SENTINEL2_L2A.name] + assert encoded.shape == tokens.shape + assert torch.isfinite(encoded).all() + + expected_time = get_1d_sincos_pos_encoding(torch.arange(5), 4) + actual_time = encoded[0, 0, 0, :, 0, 4:8] + assert torch.allclose(actual_time, expected_time, atol=1e-5) + + +def test_composite_encodings_drops_legacy_pos_embed_on_load() -> None: + """Strict state dict loading ignores the removed fixed temporal table.""" + encodings = CompositeEncodings( + embedding_size=16, + supported_modalities=[Modality.SENTINEL2_L2A], + ) + state_dict = encodings.state_dict() + state_dict["pos_embed"] = torch.zeros((3, 4)) + + encodings.load_state_dict(state_dict, strict=True) + + def test_band_dropout_disabled_by_default() -> None: """Test that Encoder leaves band dropout disabled at construction.""" encoder = Encoder( From 0aeb510a04206759ef5cf220c862a679eec45d0f Mon Sep 17 00:00:00 2001 From: "robin.cole@earthdaily.com" Date: Mon, 20 Jul 2026 13:43:44 +0100 Subject: [PATCH 2/2] fix tests --- .../olmoearth_pretrain_v1/nn/flexi_vit.py | 48 ++++++++----------- tests/test_flexi_vit.py | 20 ++++---- 2 files changed, 30 insertions(+), 38 deletions(-) diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py index 6c81ec0..136886f 100644 --- a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py +++ b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py @@ -2,7 +2,6 @@ import logging import math -import warnings from dataclasses import dataclass from enum import StrEnum from typing import Any, NamedTuple @@ -854,8 +853,8 @@ def __init__( embedding_size: Size of token embeddings supported_modalities: Which modalities from Modality this model instantiation supports - max_sequence_length: Deprecated, has no effect. Temporal position - encodings are now computed on-the-fly. + max_sequence_length: Size of the checkpoint-compatible temporal + encoding table. Longer sequences are computed on-the-fly. learnable_channel_embeddings: Whether to use learnable channel embeddings random_channel_embeddings: Initialize channel embeddings randomly (zeros if False) tokenization_config: Optional config for custom band groupings @@ -872,13 +871,8 @@ def __init__( f"position_encoding must be one of {PositionEncoding.values()}, " f"got {position_encoding}" ) - if max_sequence_length is not None: - warnings.warn( - "max_sequence_length is deprecated and has no effect. " - "Temporal position encodings are now computed on-the-fly.", - DeprecationWarning, - stacklevel=2, - ) + if max_sequence_length is None: + max_sequence_length = 12 self.embedding_size = embedding_size self.supported_modalities = supported_modalities self.supported_modality_names = [ @@ -892,10 +886,15 @@ def __init__( # we have 4 embeddings types (pos_in_time, pos_in_space, month, channel) so each get # 0.25 of the dimension self.embedding_dim_per_embedding_type = int(embedding_size * 0.25) - # Temporal position encodings are computed on-the-fly via - # get_1d_sincos_pos_encoding so that any number of timesteps is supported - # without a pre-allocated table. - self._register_load_state_dict_pre_hook(self._drop_pos_embed_hook) + # Retain the frozen table for exact compatibility with existing + # checkpoints. Positions beyond it are computed on-the-fly. + self.pos_embed = nn.Parameter( + get_1d_sincos_pos_encoding( + torch.arange(max_sequence_length), + self.embedding_dim_per_embedding_type, + ), + requires_grad=False, + ) # Month encodings month_tab = get_month_encoding_table(self.embedding_dim_per_embedding_type) self.month_embed = nn.Embedding.from_pretrained(month_tab, freeze=True) @@ -935,15 +934,6 @@ def _init_weights(self, m: nn.Module) -> None: # TODO: fix the dtype here nn.init.constant_(m.bias, 0).to(torch.float32) - @staticmethod - def _drop_pos_embed_hook( - state_dict: dict, prefix: str, *args: object, **kwargs: object - ) -> None: - """Drop legacy pos_embed from old checkpoints so strict loading succeeds.""" - key = prefix + "pos_embed" - if key in state_dict: - del state_dict[key] - @staticmethod def calculate_gsd_ratio(input_res: float, patch_size: int) -> float: """Calculate the Ground Sample Distance ratio.""" @@ -1043,11 +1033,13 @@ def _apply_encodings_per_modality( # Slot-index temporal encoding (additive). Skipped when 3D RoPE # handles temporal position rotationally inside attention. if not PositionEncoding.is_3d_rope(self.position_encoding): - # Time position encodings (computed on-the-fly for arbitrary t) - pos_embed = get_1d_sincos_pos_encoding( - torch.arange(t, device=device), - self.embedding_dim_per_embedding_type, - ) + if t <= self.pos_embed.shape[0]: + pos_embed = self.pos_embed[:t].to(device) + else: + pos_embed = get_1d_sincos_pos_encoding( + torch.arange(t, device=device), + self.embedding_dim_per_embedding_type, + ) time_embed = repeat(pos_embed, f"t d -> {ein_string}", **ein_dict) modality_embed[..., n : n * 2] += time_embed diff --git a/tests/test_flexi_vit.py b/tests/test_flexi_vit.py index 8d4f624..0c75c7a 100644 --- a/tests/test_flexi_vit.py +++ b/tests/test_flexi_vit.py @@ -1,4 +1,3 @@ -import pytest import torch from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.encodings import ( @@ -13,13 +12,12 @@ def test_composite_encodings_compute_temporal_positions_on_the_fly() -> None: """CompositeEncodings supports sequences longer than deprecated config value.""" - with pytest.warns(DeprecationWarning): - encodings = CompositeEncodings( - embedding_size=16, - supported_modalities=[Modality.SENTINEL2_L2A], - max_sequence_length=3, - ) - assert "pos_embed" not in encodings.state_dict() + encodings = CompositeEncodings( + embedding_size=16, + supported_modalities=[Modality.SENTINEL2_L2A], + max_sequence_length=3, + ) + assert encodings.pos_embed.shape == (3, 4) tokens = torch.zeros((1, 2, 2, 5, Modality.SENTINEL2_L2A.num_band_sets, 16)) timestamps = torch.tensor( @@ -41,16 +39,18 @@ def test_composite_encodings_compute_temporal_positions_on_the_fly() -> None: assert torch.allclose(actual_time, expected_time, atol=1e-5) -def test_composite_encodings_drops_legacy_pos_embed_on_load() -> None: - """Strict state dict loading ignores the removed fixed temporal table.""" +def test_composite_encodings_loads_temporal_position_table() -> None: + """Strict state dict loading restores the frozen temporal table.""" encodings = CompositeEncodings( embedding_size=16, supported_modalities=[Modality.SENTINEL2_L2A], + max_sequence_length=3, ) state_dict = encodings.state_dict() state_dict["pos_embed"] = torch.zeros((3, 4)) encodings.load_state_dict(state_dict, strict=True) + assert torch.equal(encodings.pos_embed, torch.zeros((3, 4))) def test_band_dropout_disabled_by_default() -> None: