Fixed the llama model

[yiliu30]

In the training mode (model.setup_caches(..., training=True)) with input_pos is None, the freq_cis is overridden by L208.

ao/torchao/_models/llama/model.py

Lines 198 to 208 in e2dad4a

	if input_pos is None:
	mask = None
	freqs_cis = self.freqs_cis[:idx.shape[1]]
	elif not self.linear_causal_mask:
	mask = self.causal_mask[None, None, input_pos]
	elif len(input_pos)>1 and self.linear_causal_mask: # prefill for linear causal mask
	mask = torch.tril(torch.ones(len(input_pos), self.max_seq_length, dtype=torch.bool, device=input_pos.device)).unsqueeze(0).unsqueeze(0)
	else: # decode_one_token for linear causal mask
	self.causal_mask[0,0,0,input_pos] = 1
	mask = self.causal_mask
	freqs_cis = self.freqs_cis[input_pos]

This PR attempts to fix this issue, and added some tests for the Llama model.

Test

 pytest  -sv ./test/test_ao_models.py 

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/769

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 64480e7 with merge base 05224a9 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[HDCharles]

This is relevant only to the training path that @gau-nernst added recently. He would be the one to ask, i'll add him as a reviewer.

[gau-nernst]

Sorry for the delay! From what I see, this PR fixes the situation when model is used for inference and input_pos is None i.e.

model = Transformer.from_name(...)
model.setup_caches(training=False)
model(input_ids)

However, it seems like the model is not used this way anywhere in torchao? Or do you have something else planned in your mind?

Otherwise, the change looks fine to me.

Note: from what I know, here are the 2 places where this Llama model is used

• For inference in torchao/_models/llama/generate.py, where input_pos is always passed explicitly
• For training in benchmarks/quantized_training/pretrain_llama2.py (KV cache is not initialized), where input_pos is None -> no need input_pos since we don't have KV cache.

[yiliu30]

Hi @gau-nernst, thanks for providing such detailed background. I'm using this Llama model for auto-round. The scenario is similar to the training case you mentioned above(not using KV cache and no input_pos). More details can be found here.
https://github.com/yiliu30/torchao-fork/blob/21686f1c87b2961ee0245740e2dcaa6e7fbc4f3a/torchao/_models/llama/generate.py#L245-L267

If the input_pos is None, I believe we should select self.freqs_cis[:idx.shape[1]] for freqs_cis instead of self.freqs_cis. This is what this PR fixed. I also added some UTs, though we may need more in the future to ensure that modifications to this model don't break the generation benchmark.

[gau-nernst]

self.freqs_cis[:idx.shape[1]] is already done in the latest main I think. Also evident in this PR diff.

ao/torchao/_models/llama/model.py

Lines 198 to 200 in e2dad4a

	if input_pos is None:
	mask = None
	freqs_cis = self.freqs_cis[:idx.shape[1]]

Given your use case, which also does not use KV cache, I don't think having input_pos = torch.arange(0, idx.shape[1], device=idx.device) is necessary?

I agree that having test is good to prevent regression. Since you are already working on this PR, can you also add a short docstring/comment in Llama's forward about how to use this model? i.e. for inference (w/ KV-cache + input_pos) and training (no KV-cache + no input_pos).

[gau-nernst]

This will test the case "KV cache + input_pos". If this is not required, perhaps we can test for inference and training separately? i.e.

# inference case
random_model.setup_caches(training=False)
random_model(input_ids, input_pos)  # input_pos is not None

# training case
random_model.setup_caches(training=True)
random_model(input_ids)  # no input_pos

[yiliu30]

self.freqs_cis[:idx.shape[1]] is already done in the latest main I think. Also evident in this PR diff.

ao/torchao/_models/llama/model.py

Lines 198 to 200 in e2dad4a

	if input_pos is None:
	mask = None
	freqs_cis = self.freqs_cis[:idx.shape[1]]

Yeah, but L208 update the freqs_cis again.

ao/torchao/_models/llama/model.py

Lines 198 to 208 in e2dad4a

	if input_pos is None:
	mask = None
	freqs_cis = self.freqs_cis[:idx.shape[1]]
	elif not self.linear_causal_mask:
	mask = self.causal_mask[None, None, input_pos]
	elif len(input_pos)>1 and self.linear_causal_mask: # prefill for linear causal mask
	mask = torch.tril(torch.ones(len(input_pos), self.max_seq_length, dtype=torch.bool, device=input_pos.device)).unsqueeze(0).unsqueeze(0)
	else: # decode_one_token for linear causal mask
	self.causal_mask[0,0,0,input_pos] = 1
	mask = self.causal_mask
	freqs_cis = self.freqs_cis[input_pos]

[gau-nernst]

Thank you for your reply. It makes sense to separate the logic for creating mask, input_pos, freq_cis like you did here. However, creating input_pos here (for training case) is now not needed?

[yiliu30]

I've also tested the inference case where input_pos is None. Is input_pos a required argument for model.forward in inference mode?

[gau-nernst]

inference case where input_pos is None

This is what I mentioned earlier. I think right now there is no code that uses inference w/o input_pos, and your auto-round PR doesn't seem to need it also. I think it's fine to support this case (though we will have a tiny inefficiency -> create an unneeded input_pos during training. probably insignificant). Perhaps others can have other comments.

Can you update the PR description to describe the problem this PR fixes clearer? i.e. when input_pos is None (during training), freq_cis is overridden by line xxx. I will approve the PR once you added tests for training mode (model.setup_caches(training=True)) and add a short docstring/comment in model.forward().

[yiliu30]

Got it, thanks for the clarification! I updated the PR description, UTs and docstring. Please review it again.