[Feature] Allow targeting multiples of sequential targets

[kylesayrs]

Background

LLM Compressor uses sequential onloading to onload subgraphs of a model. How and which subgraphs are created is determined by the sequential_targets argument, either passed by the user or inferred from the model definition (typically one DecoderLayer).

The choice of sequential targets comes with tradeoffs. Larger sequential targets are more runtime efficient, but use more memory. Smaller sequential targets use less memory, but are less runtime efficeint.

Currently, a user can only specify the sequential_targets argument. However, there are cases where a user will want to pack multiple sequential targets into a single subgraph (for example, onload two decoder layers per subgraph, rather than just one).

The logic for creating subgraphs from targets is implemented by the topological_partition function.

Requested Changes

1. Design an interface for allow users to target multiple sequential targets per subgraph (for example, if your GPUs can fit 2-3 decoder layers)
2. Modify topological_partition to allow multiple targets to be assigned to the same subgraph
3. Test the feature. You should see higher vram usage and lower runtime

[aayush7511]

Hi @kylesayrs can I work on this?

I've read through topological_partition and I have an idea of how to solve this. My plan is to:

1. Add a num_targets_per_subgraph: int = 1 parameter to the function arguments
2. Modify topological_partition to group N consecutive targets into the same subgraph before partitioning
3. Add a test verifying higher VRAM usage and lower runtime with N > 1

I'm happy to open a draft PR once I have a working implementation if you'd like to review the approach early. Let me know if there's a preferred interface design or anything I should be aware of before starting.

[kylesayrs]

@aayush7511 I think that sounds like a great plan! I'll assign this ticket to you

[aayush7511]

Hi @kylesayrs, the draft PR is ready. I'm currently working on tests.

[aayush7511]

Hi @kylesayrs,

I had a quick question.
#2493 (comment)

This Is my first time contributing, so I wasn't sure if I should drop the question here or chained to the PR. Feel free to respond when you have a minute.

Thanks!