feat: replace sliding window type with offset (#1989)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: k223kim

litgpt/config.py

@@ -3,7 +3,7 @@
 from copy import deepcopy
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Callable, Literal, Optional, Type, Union
+from typing import Any, Literal, Optional, Type, Union
 
 import torch
 import yaml
@@ -60,7 +60,7 @@ class Config:
     sliding_window_size: Optional[int] = None
     sliding_window_layer_placing: Optional[Literal["all", "interleaved"]] = None
     sliding_window_layer_stride: Optional[int] = None
-    sliding_window_type: Optional[Literal["gemma3"]] = None
+    sliding_window_offset: int = 0
     # if `attention_logit_softcapping` is used, cannot use optimized
     # `torch.nn.functional.scaled_dot_product_attention` (which implements
     # Flash attention), may result in higher memory and runtime footprint.
@@ -118,10 +118,7 @@ def __post_init__(self):
                 else self.sliding_window_layer_stride
             )
 
-        SLIDING_WINDOW_TYPE_TO_MAP_FN: dict[Literal["gemma3"], Callable[[int], int]] = {"gemma3": lambda x: x + 1}
-        self.sliding_window_block_idx_map_fn = (
-            lambda x: x if self.sliding_window_type is None else SLIDING_WINDOW_TYPE_TO_MAP_FN[self.sliding_window_type]
-        )
+        self.sliding_window_block_idx_map_fn = lambda x: x + self.sliding_window_offset
 
     @classmethod
     def from_name(cls, name: str, **kwargs: Any) -> Optional[Self]: