Skip to content

GitLab

Commit dff3d157 authored by Casper's avatar Casper
Browse files

Allow PytorchGELUTanh

parent de8e7ffc
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 3 deletions
awq/quantize/scale.py
View file @ dff3d157
import torch import torch
import torch.nn as nn import torch.nn as nn
from typing import Tuple from typing import Tuple
from awq.modules.act import ScaledActivation from awq.modules.act import ScaledActivation
from transformers.activations import NewGELUActivation
from awq.utils.module import get_op_by_name, set_op_by_name from awq.utils.module import get_op_by_name, set_op_by_name
from transformers.models.bloom.modeling_bloom import BloomGelu from transformers.models.bloom.modeling_bloom import BloomGelu
from transformers.models.llama.modeling_llama import LlamaRMSNorm from transformers.models.llama.modeling_llama import LlamaRMSNorm
from transformers.activations import NewGELUActivation, PytorchGELUTanh
allowed_norms = [nn.LayerNorm, LlamaRMSNorm] allowed_norms = [nn.LayerNorm, LlamaRMSNorm]
allowed_act_fns = [nn.GELU, BloomGelu, NewGELUActivation] allowed_act_fns = [nn.GELU, BloomGelu, NewGELUActivation, PytorchGELUTanh]
@torch.no_grad() @torch.no_grad()
def apply_clip(module, clip_list: Tuple[str, torch.Tensor]): def apply_clip(module, clip_list: Tuple[str, torch.Tensor]):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment