[Tracing] Allow torch.Sizes to be iterated

[kylesayrs]

Purpose

• Reduce user burden when tracing models by automatically handling the iteration of torch.Size types

The following code

input_shape = hidden_states.shape[:-1]
hidden_shape = (*input_shape, -1, self.head_dim) # iterate shape

gets graph compiled into

input_shape = getattr(hidden_states, "shape")
input_shape_2 = __getitem__(input_shape, slice(None, -1, None))
hidden_shape = (__getitem__(input_shape_2, 0), __getitem__(input_shape_2, 1), -1, head_dim)

Changes

• Catch any attempts to iterate torch.Size proxies. Because lengths are assumed to be static during tracing, we can replace the iteration with an iteration of __getitem__ calls of each of the indices

Testing

• Can now run sequential pipeline with llama3 with Linear granularity without custom code definition

MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"], sequential_targets="Linear")

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[brian-dellabetta]

neat!

[brian-dellabetta]

maybe this?

Suggested change

	# their values are dynamic
	# their values as dynamic