Upload folder using huggingface_hub

config.json

@@ -6,7 +6,6 @@
     "AutoConfig": "configuration_baichuan.BaichuanConfig",
     "AutoModelForCausalLM": "modeling_baichuan.BaichuanForCausalLM"
   },
-  "tokenizer_class": "BaichuanTokenizer",
   "bos_token_id": 1,
   "eos_token_id": 2,
   "hidden_act": "silu",

configuration_baichuan.py

@@ -46,7 +46,6 @@ class BaichuanConfig(PretrainedConfig):
         bos_token_id=1,
         eos_token_id=2,
         tie_word_embeddings=False,
-        z_loss_weight=0,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -59,7 +58,6 @@ class BaichuanConfig(PretrainedConfig):
         self.initializer_range = initializer_range
         self.rms_norm_eps = rms_norm_eps
         self.use_cache = use_cache
-        self.z_loss_weight = z_loss_weight
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,

modeling_baichuan.py

@@ -502,7 +502,6 @@ class NormHead(nn.Module):
     def forward(self, hidden_states):
         if self.training:
             norm_weight = nn.functional.normalize(self.weight)
-            self.first_flag = True
         elif self.first_flag:
             self.first_flag = False
             self.weight = nn.Parameter(nn.functional.normalize(self.weight))
@@ -529,7 +528,7 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
         self.model = BaichuanModel(config)
 
         self.lm_head = NormHead(config.hidden_size, config.vocab_size, bias=False)
-        if hasattr(config, "quantization_config") and isinstance(config.quantization_config, dict) and config.quantization_config.get('load_in_4bit', False):
+        if hasattr(config, "quantization_config") and config.quantization_config['load_in_4bit']:
             try:
                 from .quantizer import quantize_offline, init_model_weight_int4
             except ImportError:
@@ -609,23 +608,22 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
             model_file = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin')
             state_dict = torch.load(model_file, map_location="cpu") 
             model.is_quantized = True
-            
+                        
             device_map = kwargs.pop("device_map", None)
             torch_dtype = kwargs.pop("torch_dtype", None)
             
-            if device_map is not None:
-                kwargs = {"no_split_module_classes": model._no_split_modules}
-                target_dtype = CustomDtype.INT4
-                max_memory = get_balanced_memory(
-                    model,
-                    dtype=target_dtype,
-                    low_zero=(device_map == "balanced_low_0"),
-                    max_memory=None,
-                    **kwargs,
-                )
-                kwargs["max_memory"] = max_memory
-                device_map = infer_auto_device_map(model, dtype=target_dtype, **kwargs)
-                
+            kwargs = {"no_split_module_classes": model._no_split_modules}
+            target_dtype = CustomDtype.INT4
+            max_memory = get_balanced_memory(
+                model,
+                dtype=target_dtype,
+                low_zero=(device_map == "balanced_low_0"),
+                max_memory=None,
+                **kwargs,
+            )
+            kwargs["max_memory"] = max_memory
+            
+            device_map = infer_auto_device_map(model, dtype=target_dtype, **kwargs)
             model = init_model_weight_int4(config, model, state_dict)
             
             # Set model in evaluation mode to deactivate DropOut modules by default
@@ -706,11 +704,9 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
             loss_fct = CrossEntropyLoss()
             shift_logits = shift_logits.view(-1, self.config.vocab_size)
             shift_labels = shift_labels.view(-1)
-            softmax_normalizer = shift_logits.max(-1).values ** 2
-            z_loss = self.config.z_loss_weight * softmax_normalizer.mean()
             # Enable model parallelism
             shift_labels = shift_labels.to(shift_logits.device)
-            loss = loss_fct(shift_logits, shift_labels) + z_loss
+            loss = loss_fct(shift_logits, shift_labels)
 
         if not return_dict:
             output = (logits,) + outputs[1:]