Make config/class properties on ComposerMPTForCausalLM (#1227)

mosaicml · May 22, 2024 · c891bed · c891bed
1 parent 8e29698
commit c891bed
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/llmfoundry/models/mpt/modeling_mpt.py b/llmfoundry/models/mpt/modeling_mpt.py
@@ -19,6 +19,7 @@
     MutableMapping,
     Optional,
     Tuple,
+    Type,
     Union,
 )
 
@@ -1082,9 +1083,7 @@ def __init__(
 
         additional_train_metrics = additional_train_metrics or []
 
-        model = MPTForCausalLM(
-            MPTConfig(use_train_metrics=use_train_metrics, **kwargs),
-        )
+        model = self.model_class(self.config_class(**kwargs),)
 
         use_train_metrics = use_train_metrics
         train_metric_names = DEFAULT_CAUSAL_LM_TRAIN_METRICS + additional_train_metrics
@@ -1134,6 +1133,14 @@ def __init__(
                 f'Specified loss_fn={self.loss_fn} not recognized. `loss_fn` must be one of [`fused_crossentropy`, `torch_crossentropy`].',
             )
 
+    @property
+    def model_class(self) -> Type[MPTForCausalLM]:
+        return MPTForCausalLM
+
+    @property
+    def config_class(self) -> Type[MPTConfig]:
+        return MPTConfig
+
     def get_targets(self, batch: Mapping) -> torch.Tensor:
         targets = torch.roll(batch['labels'], shifts=-1)
         targets[:, -1] = -100