Skip to content

Fix Llama4 example #2846

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix Llama4 example #2846

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
5 changes: 4 additions & 1 deletion torchao/prototype/moe_quant/llama4_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,10 @@ def convert_fn(module):


model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
model = Llama4ForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)

dtype = torch.bfloat16
torch.set_default_dtype(dtype)
model = Llama4ForCausalLM.from_pretrained(model_id, torch_dtype=dtype)
tokenizer = AutoTokenizer.from_pretrained(model_id)

_replace_with_custom_fn_if_matches_filter(
Expand Down
2 changes: 1 addition & 1 deletion torchao/prototype/moe_quant/quantizable_moe_modules.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ def __init__(
def forward(self, x: Tensor) -> Tensor:
batch_size = x.shape[0]
x = x.view(-1, self.hidden_dim) # x: [T, D]
scores = self.router(x) # [T, E]
scores = self.router(x)[0] # [T, E]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The original router returns router_scores and router_logits.
https://github.com/huggingface/transformers/blob/19ffe0219dae122203f9726669f88ef1c6ea3bb4/src/transformers/models/llama4/modeling_llama4.py#L143

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

router here is a nn.Linear, not Llama4Router actually, see L21?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This rounter was replaced with the original rounter.

ao/torchao/prototype/moe_quant/llama4_quant.py

Lines 48 to 53 in 9d01b43

router = module.router
up_proj = module.experts.gate_up_proj
w1, w3 = up_proj.permute(0, 2, 1).chunk(2, dim=1)
w2 = module.experts.down_proj.permute(0, 2, 1)
new_mod.router = router

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so this module does not run by itself? seems quite confusing

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The MOEFeedForwardAOQuantizable was used, but it seems that its rounter shouldn’t be quantized in order to preserve accuracy. Could you confirm that? @HDCharles

scores = F.softmax(scores, dim=-1)
scores, expert_indices = torch.topk(
scores, self.top_k, dim=-1
Expand Down
Loading