sync base to chat

modeling_baichuan.py

@@ -528,7 +528,6 @@ 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 config.quantization_config['load_in_4bit']:
         if hasattr(config, "quantization_config") and isinstance(config.quantization_config, dict) and config.quantization_config.get('load_in_4bit', False):
             try:
                 from .quantizer import quantize_offline, init_model_weight_int4
@@ -609,22 +608,23 @@ 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)
             
-            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)
+            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)
+                
             model = init_model_weight_int4(config, model, state_dict)
             
             # Set model in evaluation mode to deactivate DropOut modules by default