Fixes

__init__.py

@@ -19,7 +19,7 @@ from .layers import (
     FlexBertUnpadPostNormLayer,
     FlexBertUnpadPreNormLayer,
 )
-from .model import (
+from .modeling_flexbert import (
     BertLMPredictionHead,
     BertModel,
     BertForMaskedLM,
@@ -68,6 +68,6 @@ __all__ = [
     "FlexBertForMaskedLM",
     "FlexBertForSequenceClassification",
     "FlexBertForMultipleChoice",
-    "IndexFirstAxis,
+    "IndexFirstAxis",
     "IndexPutFirstAxis"
 ]

attention.py

@@ -20,12 +20,15 @@ from typing import Optional
 import importlib.metadata
 import logging
 import math
-
+import sys
+import os
+# Add src folder root to path to allow us to use relative imports regardless of what directory the script is run from
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))
 import bert_padding
 from .configuration_bert import FlexBertConfig, maybe_add_padding
 from .normalization import get_norm_layer
 from .initialization import ModuleType, init_weights
-import src.utils  # noqa: F401
+import utils  # noqa: F401
 
 IMPL_USE_FLASH3 = False
 IMPL_USE_FLASH2 = False

initialization.py

@@ -14,7 +14,7 @@ from typing import Optional, Union
 import torch
 import torch.nn as nn
 
-from src.utils import StrEnum
+from utils import StrEnum
 
 from .configuration_bert import FlexBertConfig
 from .normalization import RMSNorm

modeling_flexbert.py

@@ -69,8 +69,8 @@ from transformers.models.bert.modeling_bert import BertPreTrainedModel
 
 from bert_padding import index_put_first_axis
 
-from src.bert_layers.activation import get_act_fn
-from src.bert_layers.attention import (
+from .activation import get_act_fn
+from .attention import (
     FlexBertPaddedAttention,
     FlexBertPaddedParallelAttention,
     FlexBertPaddedRopeAttention,
@@ -80,15 +80,15 @@ from src.bert_layers.attention import (
     FlexBertUnpadRopeAttention,
     FlexBertUnpadRopeParallelAttention,
 )
-from src.bert_layers.configuration_bert import FlexBertConfig
-from src.bert_layers.embeddings import (
+from .configuration_bert import FlexBertConfig
+from .embeddings import (
     BertAlibiEmbeddings,
     FlexBertAbsoluteEmbeddings,
     FlexBertCompiledSansPositionEmbeddings,
     FlexBertSansPositionEmbeddings,
     get_embedding_layer,
 )
-from src.bert_layers.initialization import (
+from .initialization import (
     ModuleType,
     TileLinear,
     TileMode,
@@ -97,7 +97,7 @@ from src.bert_layers.initialization import (
     tile_linear,
     tile_norm,
 )
-from src.bert_layers.layers import (
+from .layers import (
     BertAlibiEncoder,
     BertPooler,
     BertPredictionHeadTransform,
@@ -112,10 +112,9 @@ from src.bert_layers.layers import (
     FlexBertUnpadPreNormLayer,
     get_encoder_layer,
 )
-from src.bert_layers.loss import get_loss_fn
-from src.bert_layers.mlp import FlexBertGLU, FlexBertMLP, FlexBertParallelGLU
-from src.bert_layers.normalization import get_norm_layer
-from src.bert_layers.padding import pad_input, unpad_input
+from .mlp import FlexBertGLU, FlexBertMLP, FlexBertParallelGLU
+from .normalization import get_norm_layer
+from .padding import pad_input, unpad_input
 
 logger = logging.getLogger(__name__)
 
@@ -867,14 +866,16 @@ class FlexBertPreTrainedModel(BertPreTrainedModel):
 
     def _init_module_weights(self, module: nn.Module):
         """
-        Custom weight init of modules using src.bert_layers.initialization.init_weights
+        Custom weight init of modules using .initialization.init_weights
         Currently only supports init of embedding modules
         """
         assert isinstance(module, nn.Module)
         if isinstance(module, nn.Embedding):
             init_weights(self.config, module, type_of_module=ModuleType.emb)
         else:
-            raise NotImplementedError("Custom weight init for the given module is not supported")
+            print("Custom weight init for the given module is not supported")
+            print(module)
+            # raise NotImplementedError("Custom weight init for the given module is not supported")
 
 
 class FlexBertModel(FlexBertPreTrainedModel):
@@ -1010,8 +1011,6 @@ class FlexBertForMaskedLM(FlexBertPreTrainedModel):
             decoder_weights = nn.Linear(config.hidden_size, config.vocab_size, bias=False).weight
         self.decoder = nn.Linear(decoder_weights.size(1), decoder_weights.size(0), bias=config.decoder_bias)
         self.decoder.weight = decoder_weights
-
-        self.loss_fn = nn.CrossEntropyLoss() if not hasattr(config, "loss_function") else get_loss_fn(config)
         self.fa_ce = getattr(config, "loss_function", "cross_entropy") == "fa_cross_entropy"
         self.return_z_loss = config.loss_kwargs.get("return_z_loss", False)
         self.unpad_embeddings = config.unpad_embeddings