Skip to content

Count PackedSequence cell steps in the flops profiler RNN hook - #8507

Open
ebarkhordar wants to merge 1 commit into
deepspeedai:masterfrom
ebarkhordar:fix/4333-flops-profiler-packedsequence
Open

Count PackedSequence cell steps in the flops profiler RNN hook#8507
ebarkhordar wants to merge 1 commit into
deepspeedai:masterfrom
ebarkhordar:fix/4333-flops-profiler-packedsequence

Conversation

@ebarkhordar

Copy link
Copy Markdown
Contributor

_rnn_forward_hook reads inp.shape[0] and inp.shape[1], and nn.RNN, nn.GRU and nn.LSTM all accept a PackedSequence there, so profiling a packed-input RNN raises the AttributeError you hit.

Slight correction to the patch in the issue: inp.data is (sum(lengths), input_size), so data.shape[1] is the input size and not a sequence length. Using it multiplies the count by that factor. On a 2-layer LSTM(8, 16) with 4 sequences of length 5 it gives a multiplier of 160 where 20 is right, so the crash turns into a silent 8x overcount.

What the hook wants is the number of cell steps, and a PackedSequence runs one per row of .data. So batch_size = inp.data.shape[0] and seq_length = 1.

Added a test over one model with three inputs: the dense batch, the same batch packed at full length, and lengths [5, 3, 2, 1]. Full-length packing comes out equal to the dense count (149760), and the ragged one at 11 of its 20 steps (82368). The two packed cases fail on master. The dense case is there as the reference the other two are checked against, and the RNN hook had no test before this.

Only ran it on CPU.

Fixes #4333

_rnn_forward_hook read inp.shape[0] and inp.shape[1] off the first positional
argument, but nn.RNN, nn.GRU and nn.LSTM all accept a PackedSequence there and
it has no .shape, so profiling any packed-input RNN raised AttributeError.

A PackedSequence runs one cell step per row of .data, so .data.shape[0] is
already the batch_size * seq_length product the count multiplies by. Reading
.data.shape[1] as the sequence length would scale the result by input_size.

Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3dc3a8c7d6

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

inp = input[0]
batch_size = inp.shape[0]
seq_length = inp.shape[1]
if isinstance(inp, nn.utils.rnn.PackedSequence):

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Add the required sign-off trailer

This is a non-merge commit, but its message has no Signed-off-by: trailer, so it does not satisfy the repository's commit and CI requirements. Recreate the commit with git commit --signoff using the configured author name and email.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] FlopsProfiler cannot handle input of type torch.nn.utils.rnn.PackedSequence

1 participant