Addressed missed out comments in PR #1

pyproject.toml

@@ -58,7 +58,6 @@ include = [
     "instructlab.dolomite.hf_models.modeling_utils.normalization.layernorm",
     "instructlab.dolomite.hf_models.modeling_utils.normalization.rmsnorm",
     "instructlab.dolomite.hf_models.modeling_utils.position_embedding",
-    "instructlab.dolomite.gradient_checkpointing",
     "instructlab.dolomite.utils",
 ]
 

src/instructlab/dolomite/hf_models/__init__.py

@@ -2,8 +2,9 @@
 # Extracted from https://github.com/ibm-granite/dolomite-engine
 # ----------------------------------------------------------------
 # Local
+from .config import GPTDolomiteConfig
 from .model_conversion import export_to_huggingface, import_from_huggingface
-from .models import GPTDolomiteConfig, GPTDolomiteForCausalLM, GPTDolomiteModel
+from .models import GPTDolomiteForCausalLM, GPTDolomiteModel
 from .register_hf import register_model_classes
 
 register_model_classes()

src/instructlab/dolomite/hf_models/config.py

@@ -8,7 +8,8 @@
 from .enums import AttentionHeadType, PositionEmbeddingType
 
 
-class CommonConfig(PretrainedConfig):
+class GPTDolomiteConfig(PretrainedConfig):
+    model_type = "gpt_dolomite"
     keys_to_ignore_at_inference = ["past_key_values"]
     attribute_map = {
         "hidden_size": "n_embd",
@@ -19,6 +20,8 @@ class CommonConfig(PretrainedConfig):
 
     # NOTE: initializer range is kept for backward compatiblity
     #       but it is not used anymore
+    #     : also rope_scaling is not used anymore but kept for
+    #       same reason.
 
     def __init__(
         self,

src/instructlab/dolomite/hf_models/model_conversion/bigcode.py

@@ -8,8 +8,8 @@
 from transformers import AutoConfig, AutoTokenizer, GenerationConfig, GPTBigCodeConfig
 
 # Local
+from ..config import GPTDolomiteConfig
 from ..enums import AttentionHeadType, PositionEmbeddingType
-from ..models import GPTDolomiteConfig
 
 
 def import_from_huggingface_bigcode(

src/instructlab/dolomite/hf_models/model_conversion/llama.py

@@ -6,12 +6,12 @@
 
 # Local
 from ...utils import SafeTensorsWeightsManager, download_repo
+from ..config import GPTDolomiteConfig
 from ..enums import AttentionHeadType
 from ..modeling_utils import (
     interleave_query_key_value_tensor_for_attention,
     split_query_key_value_tensor_for_attention,
 )
-from ..models import GPTDolomiteConfig
 from ..models.gpt_dolomite import (
     interleave_up_gate_tensor_for_mlp,
     split_up_gate_tensor_for_mlp,

src/instructlab/dolomite/hf_models/modeling_utils/__init__.py

@@ -14,7 +14,5 @@
     repeat_key_value,
     split_query_key_value_tensor_for_attention,
 )
-from .embedding import Embedding
-from .linear import Linear
 from .normalization import RMSNorm, get_normalization_function
-from .position_embedding import Alibi, RoPE, YaRNScaledRoPE, apply_rotary_pos_emb
+from .position_embedding import Alibi, RoPE, apply_rotary_pos_emb

src/instructlab/dolomite/hf_models/modeling_utils/attention/__init__.py

@@ -9,7 +9,7 @@
 import torch
 
 # Local
-from ...config import CommonConfig
+from ...config import GPTDolomiteConfig
 from ...enums import AttentionHeadType
 from .base import Attention
 from .flash import FlashAttention2
@@ -48,7 +48,7 @@
 
 
 def get_attention_module(
-    config: CommonConfig,
+    config: GPTDolomiteConfig,
     causal: bool,
     attention_implementation: str,
     use_padding_free_transformer: bool,

src/instructlab/dolomite/hf_models/modeling_utils/attention/base.py

@@ -5,21 +5,21 @@
 from typing import Tuple
 
 # Third Party
+from torch.nn import Linear  # replaces ParameterizedLinear
 from transformers import DynamicCache
 import torch
 import torch.nn.functional as F
 
 # Local
-from ...config import CommonConfig
+from ...config import GPTDolomiteConfig
 from ...enums import AttentionHeadType, PositionEmbeddingType
-from ..linear import Linear
 from ..position_embedding import apply_rotary_pos_emb
 from .utils import repeat_key_value
 
 
 class Attention(torch.nn.Module):
     def __init__(
-        self, config: CommonConfig, causal: bool, layer_idx: int = None
+        self, config: GPTDolomiteConfig, causal: bool, layer_idx: int = None
     ) -> None:
         super().__init__()
 

src/instructlab/dolomite/hf_models/modeling_utils/normalization/__init__.py

@@ -5,8 +5,7 @@
 import torch
 
 # Local
-from .layernorm import get_layernorm
-from .rmsnorm import RMSNorm, get_rmsnorm
+from .norms import RMSNorm, get_layernorm, get_rmsnorm
 
 _NORMALIZATION_FUNCTIONS = {
     "layernorm": get_layernorm,

src/instructlab/dolomite/hf_models/modeling_utils/normalization/norms.py

@@ -0,0 +1,81 @@
+# ----------------------------------------------------------------
+# Extracted from https://github.com/ibm-granite/dolomite-engine
+# ----------------------------------------------------------------
+
+# Standard
+import numbers
+
+# Third Party
+import torch
+
+# ---------------- LayerNorm ---------------
+
+_LAYERNORM_MODULES = {
+    "torch": torch.nn.LayerNorm,
+}
+
+
+def get_layernorm(
+    normalized_shape: int,
+    eps: float,
+    normalization_implementation: str = "torch",
+) -> torch.nn.LayerNorm:
+    if normalization_implementation in _LAYERNORM_MODULES:
+        return _LAYERNORM_MODULES[normalization_implementation](
+            normalized_shape=normalized_shape, eps=eps
+        )
+
+    raise ValueError(
+        f"unexpected `normalization_implementation` {normalization_implementation}"
+    )
+
+
+# --------------- RMS Norm ---------------
+# ----------------------------------------------------------------
+# Extracted from https://github.com/ibm-granite/dolomite-engine
+# ----------------------------------------------------------------
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6) -> None:
+        super().__init__()
+
+        self.weight = torch.nn.Parameter(torch.ones(normalized_shape))
+        self.eps = eps
+
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        self.normalized_shape = normalized_shape
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        input_dtype = input.dtype
+
+        input = input.to(torch.float32)
+        variance = input.pow(2).mean(-1, keepdim=True)
+        input = input * torch.rsqrt(variance + self.eps)
+
+        return self.weight * input.to(input_dtype)
+
+    def extra_repr(self) -> str:
+        return f"{self.normalized_shape}, eps={self.eps}"
+
+    def reset_parameters(self) -> None:
+        torch.nn.init.ones_(self.weight)
+
+
+_RMSNORM_MODULES = {"torch": RMSNorm}
+
+
+def get_rmsnorm(
+    normalized_shape: int,
+    eps: float,
+    normalization_implementation: str = "torch",
+) -> torch.nn.LayerNorm:
+    if normalization_implementation in _RMSNORM_MODULES:
+        return _RMSNORM_MODULES[normalization_implementation](
+            normalized_shape=normalized_shape, eps=eps
+        )
+
+    raise ValueError(
+        f"unexpected `normalization_implementation` {normalization_implementation}"
+    )

src/instructlab/dolomite/hf_models/modeling_utils/position_embedding/__init__.py

@@ -3,4 +3,4 @@
 # ----------------------------------------------------------------
 # Local
 from .alibi import Alibi
-from .rope import RoPE, YaRNScaledRoPE, apply_rotary_pos_emb
+from .rope import RoPE, apply_rotary_pos_emb