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/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,8 +9,8 @@
 import torch
 
 # Local
-from ...config import CommonConfig
 from ...enums import AttentionHeadType
+from ...models.gpt_dolomite.config import GPTDolomiteConfig
 from .base import Attention
 from .flash import FlashAttention2
 from .padding_free import PaddingFreeAttention
@@ -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 ...enums import AttentionHeadType, PositionEmbeddingType
-from ..linear import Linear
+from ...models.gpt_dolomite.config import GPTDolomiteConfig
 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,