Efficient way to transfer data between modules

Hello. I’d like to know efficient implementation way situation like below.

Model is consist of independent sub-modules, and there’s forwarding order between modules so previous module’s output will send as next module’s input.(this behavior is well-defined in model’s forward function.)
I’d like to add the following functionality to model:

Designate specific modules(by configuration file) and each designated module generates (huge) data that can be shared by all modules after designated module(in forwarding order)

For example, assume model has [m1, m2, m3, m4, m5] in sequence.
If I designate m2 and m4, and each generates data2 and data4. data2 can be shared in [m2, m3, m4, m5] and data4 can be shared [m4, m5].

As I mentioned above, each generated data is quite huge and shareable so I don’t want to forward this data just as part of output. (also modifying well-defined original model’s forward function is annoying job)

Currently, I’ve implemented like below:

  1. Assign empty buffer for each modules in model.
  2. Register forward hook at designated module which send generated data to following modules.

But the performance seems bad: takes 2x much more times measured same condition platform compared to original model.

The problem is each module’s does not know other module’s existence. I’ve thought about idea ‘melting’ all modules into one ‘unified’ model make the problem easier, but the model is complicated and have so many advantages that I cannot give up when keep module-based architecture.

Any ideas will be welcomed. Thanks for reading!