Fix FP32→BF16 redundant allocation during model init (#228)

* Fix FP32→BF16 redundant allocation during model init

* Improve comments in from_config method

Author: ggg-s

scripts/train_eagle3_online.py

@@ -293,18 +293,19 @@ def main():
     if draft_model_last_checkpoint:
         draft_model = (
             AutoEagle3DraftModel.from_pretrained(
-                draft_model_last_checkpoint, attention_backend=args.attention_backend
+                draft_model_last_checkpoint, attention_backend=args.attention_backend,
+                torch_dtype=torch.bfloat16
             )
             .cuda()
-            .to(torch.bfloat16)
+            
         )
     else:
         draft_model = (
             AutoEagle3DraftModel.from_config(
-                draft_model_config, attention_backend=args.attention_backend
+                draft_model_config, attention_backend=args.attention_backend,
+                torch_dtype=torch.bfloat16
             )
             .cuda()
-            .to(torch.bfloat16)
         )
     draft_model.load_embedding(args.target_model_path, embedding_key=args.embedding_key)
     draft_model.freeze_embedding()

specforge/modeling/auto.py

@@ -39,7 +39,7 @@ class AutoEagle3DraftModel(AutoModelForCausalLMBase):
     }
 
     @classmethod
-    def from_config(cls, config: PretrainedConfig, **config_kwargs):
+    def from_config(cls, config: PretrainedConfig, torch_dtype=None, **config_kwargs):
         """
         This class method takes a configuration object and create its model based on the
         _model_mapping class variable.
@@ -52,7 +52,12 @@ def from_config(cls, config: PretrainedConfig, **config_kwargs):
         """
         # get the model class from the
         _model_cls = cls._model_mapping[type(config)]
-        return _model_cls(config, **config_kwargs)
+        model = _model_cls(config, **config_kwargs)
+
+        # Convert model to specified dtype if provided
+        if torch_dtype is not None:
+            model = model.to(dtype=torch_dtype)
+        return model
 
     @classmethod
     def from_pretrained(