Auto example input generation

[glistening]

It would be nice if we don't need to write code for example inputs.

For example, we need something like:

TICO/test/modules/model/LlamaWithKVCache/model.py

Lines 36 to 70 in a3ed23e

	def get_example_inputs(self):
	past_seq_len = 511
	cur_seq_len = 1
	input_ids = torch.tensor([[812]]).to(torch.long)
	attention_mask = torch.ones(1, past_seq_len + cur_seq_len)
	position_ids = torch.tensor([[past_seq_len]]).to(torch.long)
	past_key_values = DynamicCache()
	for layer_id in range(self.config.num_hidden_layers):
	past_key_values.update(
	torch.randn(
	[
	1,
	self.config.num_attention_heads,
	past_seq_len,
	self.config.head_dim,
	]
	),
	torch.randn(
	[
	1,
	self.config.num_attention_heads,
	past_seq_len,
	self.config.head_dim,
	]
	),
	layer_id,
	)
	return (
	input_ids,
	attention_mask,
	position_ids,
	past_key_values,
	)

To correctly working example input, we need to understand much about the target model.

I tried to capture the input from user-level inputs.

I succeed to write a working version, tested with Maykeye/TinyLLama-v0.

It confirmed it works for transformers version below:

• 4.49.0 ❌ DynamicCache is not pytree-flattenable
• 4.51.3 ⭕
• 4.52.4 ⭕

A good news is that PyTorch supports transformers DynamicCahce as pytree-flatenable since 4.50.0. (though it seems some bug on MacOS, all versions after 4.50.1 would work.

I hope and guess it will work for other models by modifying model name and user inputs.

[mhs4670go]

It would be nice to have this feature:) Doesn't this have anything to do with #63? Would be helpful to know how to resolve this issue in advance.

Plus, FYI, decorator style helper is being under development as #63 (comment) said.

[glistening]

@mhs4670go Wow, I didn't know the issue. Let me close this issue, and move my comment there.

[glistening]

After reading #64, I realized it is different issue.

#64 says,

When exporting only a part of your PyTorch model (like a single layer) with torch.export

However, I don't want to export only a part, but the whole module.

The top module (for my case, LlamaForCausalLM) itself is used for both prefill and decode.

#208 generates 2 whole model for prefill and decode respectively.

$ python test/modules/model/LlamaCapturedInput/model.py

$ ls *.circle
-rw-rw-r-- 1 gyu gyu 18612884 Jul 10 10:34 decode.circle
-rw-rw-r-- 1 gyu gyu 18642160 Jul 10 10:34 prefill.circle

generate prepares the appropriate inputs and control the logic for prefill and decode.

It causes torch.export make different traces.

Perhaps #63 needs more works. However, for my case, it works w/o waiting for something.

If someone wants to get the whole prefill and decode circle for llama, use #208.

[glistening]

I confirmed it works for subgraph (LlamaDecoderLayer) with a few line modification. (It takes just a few minutes.)

See #215.

[glistening]

I added a context manager for recording forward's inputs.

Just run generate() with recording enabled:

with torch.no_grad(), RecordingInput(model, input_to_remove=input_to_remove) as rec:
    outputs = model.generate(
        **inputs,
        max_new_tokens=32,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )
    captured_input = rec.captured_input

• No change in model.generate(...)
• Just run generate(...) in RecordingInput context

It will capture model.forward(...)'s inputs to captured_input.

Then, just pass this captured input as input for tico.convert.

circle_model = tico.convert(model, captured_input)

[glistening]

I checked it works with Llama:

• prefill
• decode
  • 1 decode layer
  • all decoder layers

You can see the code in #217 (prefill.py, layer.py and layers.py).

I will further check with Florence2 (including fusing standard attention).

[dayo09]

After torch upgrade, kwargs validation is added to the export logic.
Due to that, changing forward function has become invalid. 😭 (happened at torch 2.7.1+cu126)

[glistening]

@dayo09 I am already using torch 2.7.1 and I have no problem. Could you please explain in detail?