Uploading optimized model files

block_config.py

@@ -0,0 +1,118 @@
+import dataclasses
+import json
+import warnings
+from dataclasses import dataclass, MISSING
+from functools import partial
+from typing import Optional, Any
+
+
+@partial(dataclass, frozen=True, kw_only=True)
+class JsonComparable:
+    def to_json(self) -> str:
+        return json.dumps(dataclasses.asdict(self))
+
+    def __eq__(self, other: "JsonComparable") -> bool:
+        return self.to_json() == other.to_json()
+
+    def __hash__(self) -> int:
+        return hash(self.to_json())
+
+    def __lt__(self, other: "JsonComparable") -> bool:
+        return self.to_json() < other.to_json()
+
+
+@partial(dataclass, frozen=True, kw_only=True)
+class SubblockConfig(JsonComparable):
+    no_op: bool = False
+    replace_with_linear: bool = False
+    sparsify: Optional[list[str]] = None
+
+    def __post_init__(self):
+        assert not (self.no_op and self.replace_with_linear)
+
+    def _force_setattr(self, name: str, value: Any) -> None:
+        """
+        Set an attribute even in frozen dataclasses.
+        Use only inside __post_init__!
+        """
+        object.__setattr__(self, name, value)
+
+
+@partial(dataclass, frozen=True, kw_only=True)
+class AttentionConfig(SubblockConfig):
+    n_heads_in_group: Optional[int] = None
+    window_length: Optional[int] = None
+    num_sink_tokens: Optional[int] = None
+    use_prefill_window_in_sink_attention: bool = False
+    unshifted_sink: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert not (self.no_op and self.replace_with_linear)
+
+        if self.no_op or self.replace_with_linear:
+            for irrelevant_att in ["n_heads_in_group", "window_length", "num_sink_tokens"]:
+                self._force_setattr(irrelevant_att, None)
+        else:
+            assert self.n_heads_in_group is not None
+
+        if self.is_sink:
+            assert not (self.unshifted_sink and self.use_prefill_window_in_sink_attention), \
+                ("Unshifted sink uses its own kind of explicit masking, not standard window. "
+                 "Set use_prefill_window_in_sink_attention to False.")
+            assert not (self.num_sink_tokens == 0 and not self.unshifted_sink), \
+                "Fake sink attention with 0 sink tokens is only supported with unshifted_sink=True"
+
+    @property
+    def prefill_sliding_window(self) -> Optional[int]:
+        if self.window_length is not None:
+            if not self.is_sink or self.use_prefill_window_in_sink_attention:
+                return self.window_length
+        return None
+
+    @property
+    def is_sliding(self) -> bool:
+        return self.prefill_sliding_window is not None
+
+    @property
+    def is_sink(self) -> bool:
+        return (
+                (self.window_length is not None)
+                and
+                (self.num_sink_tokens is not None)
+        )
+
+
+@partial(dataclass, frozen=True, kw_only=True)
+class FFNConfig(SubblockConfig):
+    ffn_mult: Optional[float] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.no_op or self.replace_with_linear:
+            self._force_setattr("ffn_mult", None)
+        else:
+            assert self.ffn_mult is not None
+            self._force_setattr("ffn_mult", round(self.ffn_mult, 6))
+
+
+@partial(dataclass, frozen=True, kw_only=True)
+class BlockConfig(JsonComparable):
+    attention: AttentionConfig = MISSING
+    ffn: FFNConfig = MISSING
+
+    def __post_init__(self):
+        """
+        Init subblock dataclasses from dicts
+        """
+        for subblock_name in dataclasses.fields(self):
+            subblock_config = getattr(self, subblock_name.name)
+            if isinstance(subblock_config, dict):
+                subblock_fields = [field.name for field in dataclasses.fields(subblock_name.type)]
+                unsupported_fields = [field_name for field_name in subblock_config.keys()
+                                      if field_name not in subblock_fields]
+                if len(unsupported_fields) > 0:
+                    warnings.warn(f"Removed unsupported fields {unsupported_fields} from {subblock_name.type.__name__}")
+                subblock_config = {k: v for k, v in subblock_config.items() if k not in unsupported_fields}
+                object.__setattr__(self, subblock_name.name,
+                                   subblock_name.type(**subblock_config))  # __setattr__ to overcome frozen=True