Skip to content

Add torchao quant (int4/int8/fp8) to llama models #1341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merrymercy merged 8 commits into from
Sep 9, 2024
Merged

Add torchao quant (int4/int8/fp8) to llama models #1341

Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0b095c8
Add torchao quant to sgl llama model for testing
jerryzh168 Sep 6, 2024
b3696b5
add torchao-config
jerryzh168 Sep 7, 2024
281c19b
add fp8
jerryzh168 Sep 7, 2024
01cc04e
Merge branch 'main' into enable-torchao
merrymercy Sep 9, 2024
6712230
lint
merrymercy Sep 9, 2024
8ae193b
move utils to torchao_utils.py
merrymercy Sep 9, 2024
69d2bb5
update
merrymercy Sep 9, 2024
b76dc72
update test cases
merrymercy Sep 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 23 additions & 2 deletions python/sglang/srt/models/llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.sampler import Sampler
from sglang.srt.model_executor.forward_batch_info import InputMetadata
ENABLE_TORCHAO = True


class LlamaMLP(nn.Module):
Expand Down Expand Up @@ -350,8 +351,17 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
if name.endswith(".bias") and name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
if hasattr(param, "weight_loader"):
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
if ENABLE_TORCHAO and name.endswith("proj.weight") and param.ndim == 2:
from torchao.quantization import quantize_, int4_weight_only
dummy_linear = torch.nn.Linear(param.shape[1], param.shape[0], bias=False)
dummy_linear.weight.data = param.data
quantize_(dummy_linear, int4_weight_only())
params_dict[name] = dummy_linear.weight

# TODO
break
else:
# Skip loading extra bias for GPTQ models.
Expand All @@ -361,6 +371,17 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)

if ENABLE_TORCHAO and name.endswith("proj.weight") and param.ndim == 2:
from torchao.quantization import quantize_, int4_weight_only
dummy_linear = torch.nn.Linear(param.shape[1], param.shape[0], bias=False)
dummy_linear.weight.data = param.data
quantize_(dummy_linear, int4_weight_only())
params_dict[name] = dummy_linear.weight

self.load_state_dict(params_dict, assign=True)




class Phi3ForCausalLM(LlamaForCausalLM):
pass
Expand Down