Sync modeling/config to silx-ai @ main (batch3)

configuration_quasar.py

@@ -2,10 +2,10 @@
 
 """
 
-from transformers.configuration_utils import PreTrainedConfig
+from transformers.configuration_utils import PretrainedConfig
 
 
-class QuasarConfig(PreTrainedConfig):
+class QuasarConfig(PretrainedConfig):
     model_type = "quasar"
     keys_to_ignore_at_inference = ["past_key_values"]
 

modeling_quasar.py

@@ -3,6 +3,7 @@
 """
 
 import math
+import os
 from dataclasses import dataclass
 
 import torch
@@ -908,6 +909,15 @@ class QuasarForCausalLM(QuasarPreTrainedModel, FLAGenerationMixin):
         self.model = QuasarModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.model.lm_head = self.lm_head
+
+        def _remap_lm_head_state_dict(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+            checkpoint_key = prefix + "model.lm_head.weight"
+            module_key = prefix + "lm_head.weight"
+            if checkpoint_key in state_dict and module_key not in state_dict:
+                state_dict[module_key] = state_dict[checkpoint_key]
+
+        self.register_load_state_dict_pre_hook(_remap_lm_head_state_dict)
         self.post_init()
 
     def get_input_embeddings(self):
@@ -967,7 +977,7 @@ class QuasarForCausalLM(QuasarPreTrainedModel, FLAGenerationMixin):
             if mask.any():
                 active_hidden = flat_hidden[mask]
                 active_labels = flat_labels[mask]
-                chunk_size = 256
+                chunk_size = int(os.environ.get("QUASAR_LM_HEAD_CHUNK_SIZE", "2048"))
                 total_loss = 0.0
                 total_tokens = active_labels.numel()
                 for i in range(0, total_tokens, chunk_size):