I came across some way to change the GPU for component modules of a larger model in the following link:
https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html 1
However, my model runs in a distributed environment (ddp) and have several sub-modules holding all of which for a forward-backward cycle each of which will decrease the size of batch I can fit for an update. But for each forward/backward, not all of these are active and therefore are not required to stay in GPU (as in I can load these for a single-batch) and otherwise leave it in the non-gpu memory.
In other words, I have parameters Theta + theta[t] (for t=1…T), where t is a particular task. I want to only load a single theta[t] for a forward and backward pass into the GPU and fit larger batches. Currently I’m holding all theta[t] in the GPU.
Is it possible to use the same semantics if it’s the same (sub)-module (theta[t]) to achieve the intention described above?