[REQUEST] How to use `Flops Profiler` to test model.generate(), with customized forward context manager

I have read the [REQUEST] How to use Flops Profiler to test model.generate() #2514 where the solution for using Flops Profiler to measure model.generate is illustrated as follows:

from transformers import T5Tokenizer, T5ForConditionalGeneration

tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5ForConditionalGeneration.from_pretrained("t5-small")

input_ids = tokenizer("translate English to German: The house is wonderful.", return_tensors="pt").input_ids

prof = FlopsProfiler(model)
# start proifle
prof.start_profile()

outputs = model.generate(input_ids)

# stop proifle and collect the profiled results
prof.stop_profile()
flops = prof.get_total_flops()
macs = prof.get_total_macs()
params = prof.get_total_params()
prof.print_model_profile()
prof.end_profile()

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

I'm currently dealing with a more complex algorithm, specially InfLLM. InfLLM designs a context manager for each layer. During each forward process, the input embeddings for the current layer to LLMs will be re-computed before calling LLMs to generate embeddings for the next layer.

Shortly speaking, InfLLM's design of context memory as well as some customized mechanisms introduce additional complexity beyond model's forward pass. I wanna try deepspeed's flops profiler to get the summarization of FLOPs for each components in the proposed framework, not only for the modules involved within the backbone model, but also for the customized components like context manager.

I will appreciate any suggestion, thanks a lot.

Start with DeepSpeed's FlopsProfiler usage around model.generate(), then inspect the InfLLM context manager referenced in inf_llm/attention/context_manager.py. Determine whether profiling can include the customized context-manager components as well as backbone modules. The issue provides no concrete acceptance criteria, so done would need to be defined as a documented or supported profiling approach.