[RFC] Lift freqs_cis as an input of models #1797

fegin · 2025-10-04T00:28:00Z

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.

[ghstack-poisoned]

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

wwwjn · 2025-10-05T03:57:27Z

torchtitan/protocols/model.py

+        self,
+        batch_size: int,
+        seq_len: int,
+    ) -> tuple[tuple[torch.Tensor, ...], tuple[int, ...]]:


Can we add some notes here what does this 2 return values mean? Seems like the first return value is the buffer itself

tianyu-l

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

tianyu-l · 2025-10-05T21:16:22Z

torchtitan/protocols/model.py

            "This model does not support attention masking/Flex Attention."
        )
+
+    def get_order_sensitive_buffers(


Naming is a bit vague. I think here we are only targeting "sequence dim" order-sensitive buffers, not the batch dim.

tianyu-l · 2025-10-05T21:24:45Z

torchtitan/models/llama3/model/model.py

+        batch_size: int,
+        seq_len: int,
+    ) -> tuple[tuple[torch.Tensor, ...], tuple[int, ...]]:
+        freqs_cis = self.freqs_cis[:seq_len].repeat(batch_size, 1, 1)


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_cis in model code

then copy it to freqs_cis, which technically is outside the model

we then send freqs_cis into model

Would like to hear your thoughts.

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.

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.

tianyu-l

SGTM, may need more benchmarking to verify the value of load balancing, which is the motivation of this change.

tianyu-l · 2025-10-07T07:12:47Z

torchtitan/models/llama3/model/model.py

+        batch_size: int,
+        seq_len: int,
+    ) -> tuple[tuple[torch.Tensor, ...], tuple[int, ...]]:
+        freqs_cis = self.freqs_cis[:seq_len].repeat(batch_size, 1, 1)


In the (maybe unlikely) case when both seq len and batch size are large, this code is creating a big buffer (which may not peak the memory anyway?).
Might be OK as in reality for ultra-long seq len we should probably set batch size = 1.

[ghstack-poisoned]

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: 8589f2f Pull-Request-resolved: #1797

[ghstack-poisoned]

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: 9d86dcc Pull-Request-resolved: #1797

[ghstack-poisoned]

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: 3e3c685 Pull-Request-resolved: #1797

[ghstack-poisoned]

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: 2d88844 Pull-Request-resolved: #1797

[ghstack-poisoned]

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: c30c532 Pull-Request-resolved: #1797

[ghstack-poisoned]

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: d6ff071 Pull-Request-resolved: #1797

[ghstack-poisoned]

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: faddc3a Pull-Request-resolved: #1797

…1776) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0) (oldest at bottom): * #1797 * __->__ #1776 **Status** 1. Change all models, including the experimental ones. 2. E2E loss verification. 3. We should add an unittest for attention. But since we don't have GPU unittest, this can be done in a separate PR. **Summary** This PR aims to refactor how TorchTitan build the attention masks and pass to model. Before this PR, init_attention_masks() is called in Trainer but the masks are stored as a class variable of FlexAttentionWrapper(). We chose this shortcut to support the case where a single model requires multiple masks. The previous design has several issues, one particular one is #1723. pytorch/pytorch#164111 proves that we can let PP split BlockMask, this PR performs the refactor to pass masks as an argument of model.forward(). The new design: 1. Model needs to provide `get_attention_masks()` that accepts `create_mask_fn`, `batch`, and `eos_id`. If the attention op is SDPA, then this API should return None as SDPA currently doesn't support varlen. But once it does, we may have to return some tuple of int that represents the mask. Justification: attention logic is technically a part of the model, but requires some information from trainer/dataloader. So it's model author's responsibility to provide some API that let trainer calls to get the masks. 2. `get_attention_masks()` will be called from the trainer and the resulting masks are passed to the model.forward(). Justification: this will allow us to fix #1723 with pytorch/pytorch#164111 and this PR. 3. Now SDPA and FlexAttention are wrapped in two different classes. ~~Note: we still have two very very thin op wrappers that are used for CP. I keep these two for the CP education purpose. But this certainly can be confusion for Titan's users. I'm opnn to merge them to AttentionOp.~~ See the discussion in #1723. **Verification** *llama3* ``` ./loss_compare.sh main 9dc16675b272ffdc3ed616e3244bcf7dc2d257f2 --steps=100 --no-seed-checkpoint --config="./torchtitan/models/llama3/train_configs/debug_model.toml" ``` *llama3 flex* ``` ./loss_compare.sh main 9dc16675b272ffdc3ed616e3244bcf7dc2d257f2 --steps=100 --no-seed-checkpoint --config="./torchtitan/models/llama3/train_configs/debug_model.toml" --baseline-train-options="--model.flavor=debugmodel_flex_attn" ``` *llama4* ``` ./loss_compare.sh main 9dc16675b272ffdc3ed616e3244bcf7dc2d257f2 --steps=100 --no-seed-checkpoint ``` *llama4 irope* ``` ./loss_compare.sh main 9dc16675b272ffdc3ed616e3244bcf7dc2d257f2 --steps=100 --no-seed-checkpoint ``` *deepseek* ``` ./loss_compare.sh main 9dc16675b272ffdc3ed616e3244bcf7dc2d257f2 --steps=100 --no-seed-checkpoint --config="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" ``` *deepseek flex* ``` ./loss_compare.sh main 9dc16675b272ffdc3ed616e3244bcf7dc2d257f2 --steps=100 --no-seed-checkpoint --config="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" --baseline-train-options="--model.flavor=debugmodel_flex_attn" ```

[ghstack-poisoned]

Update

8a42c0c

[ghstack-poisoned]

fegin requested review from tianyu-l, wconstab and wwwjn as code owners October 4, 2025 00:28

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 4, 2025

fegin mentioned this pull request Oct 4, 2025

Refactor attention and make attention mask an argument to the model #1776

Merged

wwwjn reviewed Oct 5, 2025

View reviewed changes

tianyu-l reviewed Oct 5, 2025

View reviewed changes

tianyu-l mentioned this pull request Oct 5, 2025

gpt-oss model enablement #1754

Open

tianyu-l reviewed Oct 7, 2025

View reviewed changes

Update

d6d06d5

[ghstack-poisoned]

Update

705eebd

[ghstack-poisoned]

Update

e9c4cf7

[ghstack-poisoned]

Update

d60c3d6

[ghstack-poisoned]

Update

827ba00

[ghstack-poisoned]

Update

f0dfa69

[ghstack-poisoned]

Update

b05f40d

[ghstack-poisoned]

Update

5b813a2

[ghstack-poisoned]

fegin mentioned this pull request Oct 12, 2025

[RFC][WIP][CP] Enable FlexAttention CP for llama3 #1857

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[RFC] Lift freqs_cis as an input of models #1797

[RFC] Lift freqs_cis as an input of models #1797

fegin commented Oct 4, 2025 •

edited

Loading

Uh oh!

wwwjn Oct 5, 2025

Uh oh!

tianyu-l left a comment

Uh oh!

tianyu-l Oct 5, 2025

Uh oh!

tianyu-l Oct 5, 2025

Uh oh!

fegin Oct 6, 2025 •

edited

Loading

Uh oh!

tianyu-l Oct 6, 2025

Uh oh!

tianyu-l left a comment

Uh oh!

tianyu-l Oct 7, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[RFC] Lift freqs_cis as an input of models #1797

Are you sure you want to change the base?

[RFC] Lift freqs_cis as an input of models #1797

Conversation

fegin commented Oct 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wwwjn Oct 5, 2025

Choose a reason for hiding this comment

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

tianyu-l Oct 5, 2025

Choose a reason for hiding this comment

Uh oh!

tianyu-l Oct 5, 2025

Choose a reason for hiding this comment

Uh oh!

fegin Oct 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tianyu-l Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

tianyu-l Oct 7, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

fegin commented Oct 4, 2025 •

edited

Loading

fegin Oct 6, 2025 •

edited

Loading