Merge pull request #23 from casper-hansen/yarn

YaRN support for LLaMa models

awq/models/base.py

@@ -321,7 +321,13 @@ class BaseAWQForCausalLM(nn.Module):
             
             # Load model weights
             model = AutoModelForCausalLM.from_pretrained(
-                model_filename, device_map=device_map, offload_folder="offload", offload_state_dict=True, torch_dtype=torch_dtype, use_safetensors=safetensors
+                model_filename, 
+                device_map=device_map, 
+                trust_remote_code=trust_remote_code, 
+                offload_folder="offload", 
+                offload_state_dict=True, 
+                torch_dtype=torch_dtype, 
+                use_safetensors=safetensors
             )
             model.eval()
 

awq/quantize/auto_scale.py

@@ -193,12 +193,16 @@ def apply_scale(module, scales_list, input_feat_dict=None):
         if isinstance(prev_op, nn.Linear):
             assert len(layers) == 1
             scale_fc_fc(prev_op, layers[0], scales)
-        elif isinstance(prev_op, (nn.LayerNorm, LlamaRMSNorm)):
+
+        elif any(isinstance(prev_op,t) for t in [nn.LayerNorm, LlamaRMSNorm]) \
+             or 'rmsnorm' in str(prev_op.__class__).lower():
             scale_ln_fcs(prev_op, layers, scales)
+
         elif any(isinstance(prev_op,t) for t in [nn.GELU, BloomGelu, NewGELUActivation]):
             new_module = ScaledActivation(prev_op, scales)
             set_op_by_name(module, prev_op_name, new_module)
             scale_gelu_fc(prev_op, layers[0], scales)
+            
         else:
             raise NotImplementedError(
                 f"prev_op {type(prev_op)} not supported yet!")