Skip to content

Commit 098d787

Browse files
laithsakkafacebook-github-bot
authored andcommitted
update lama export DS specs to be more accurate.
Summary: this PR pytorch/pytorch#164075 enhance value range analysis discovering a problem in the current upper bounds for DS specs for lama full DS logs if someone want to understand why in details https://www.internalfb.com/phabricator/paste/view/P1973006378 This blocks landing the PR above. Differential Revision: D83708583
1 parent d4f208d commit 098d787

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

extension/llm/export/builder.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -142,9 +142,14 @@ def __init__(
142142
{1: torch.export.Dim("token_dim", max=self.max_seq_len - 1)},
143143
)
144144
else:
145-
# Two input arguments: tokens and input_pos but input_pos is static shape
145+
# Two input arguments: tokens and input_pos but input_pos is static shape.
146+
147+
# A runtime assertion is added by torch.ops.llama.update_cache requires that
148+
# L['tokens'].size()[1] + input_pos[0].item() < self.max_seq_len
149+
# This consttaint L['tokens'].size()[1] to be elf.max_seq_len-1
150+
# run with TORCH_LOGS=+dynamic for details
146151
self.dynamic_shapes = (
147-
{1: torch.export.Dim("token_dim", max=self.max_seq_len)},
152+
{1: torch.export.Dim("token_dim", max=self.max_seq_len-1)},
148153
{"input_pos": {0: 1}},
149154
)
150155

0 commit comments

Comments
 (0)