[RFC] Lift freqs_cis as an input of models #1797
Conversation
meta-cla
bot
added
the
CLA Signed
freqs_cis is sensitive to the sequence order. CP load balancing will shuffle the samples, so each batch will have different orders. As a result, we will have to lift these order senstive buffer to the inputs and broadcast them along the batch dimension so that PP will correctly shard freqs_cis without messing up the correctness. ghstack-source-id: 49e4ec0 Pull-Request-resolved: #1797
| self, | ||
| batch_size: int, | ||
| seq_len: int, | ||
| ) -> tuple[tuple[torch.Tensor, ...], tuple[int, ...]]: |
There was a problem hiding this comment.
Can we add some notes here what does this 2 return values mean? Seems like the first return value is the buffer itself
There was a problem hiding this comment.
Since this is torchtitan-specific way of defining models, would be good if we can add a section in https://github.com/pytorch/torchtitan/blob/main/docs/composability.md
| "This model does not support attention masking/Flex Attention." | ||
| ) | ||
|
|
||
| def get_order_sensitive_buffers( |
There was a problem hiding this comment.
Naming is a bit vague. I think here we are only targeting "sequence dim" order-sensitive buffers, not the batch dim.
| batch_size: int, | ||
| seq_len: int, | ||
| ) -> tuple[tuple[torch.Tensor, ...], tuple[int, ...]]: | ||
| freqs_cis = self.freqs_cis[:seq_len].repeat(batch_size, 1, 1) |
There was a problem hiding this comment.
I wonder what's the benefit of keeping self.freqs_cis.
If seq len changes from iteration to iteration (e.g. in forge), it might be good to keep a central self.freqs_cis instead of computing it each iteration. The other benefit is that we may not want torchtitan model definition to deviate from "original" / "conventional" model definitions too much.
On the other hand, the dependency sounds indirect and error-prone:
- we create
self.freqs_cisin model code - then copy it to
freqs_cis, which technically is outside the model - we then send
freqs_cisinto model
Would like to hear your thoughts.
There was a problem hiding this comment.
Re-computation is the main reason why I decided to keep self.freqs_cis. I agree it's a bit awkward.
One alternative is to sill keep self.freqs_cis but set it as an optional field (self.freqs_cis: torch.Tensor | None) for bookkeeping only. And we only initialize it in this function. So the creation logic flow (precompute and slicing) is mainly in this function. The model code still provides precompute function. So this way we do not change the code structure too much while keeping the logic together. Not a perfect solution though.
There was a problem hiding this comment.
I think it's fine to keep the current way for now, as it sounds more lightweight change, and as I mentioned downstream application (e.g. forge, and simple generation) may change seq_len from iteration to iteration, where we can avoid recomputation this way.
Stack from ghstack (oldest at bottom):
freqs_cis is sensitive to the sequence order. CP load balancing will shuffle the samples, so each batch will have different orders. As a result, we will have to lift these order sensitive buffer to the inputs and broadcast them along the batch dimension so that PP will correctly shard freqs_cis without messing up the correctness.
Next step: once the design is finalized and people are okay with the change, we will extend this change to all models.