How to colocate generator and actor on the same device

[shendiaomo]

📚 The doc issue or request

For small models, my understanding is that torchforge supports the following colocation strategy:

1. offload torchtitan's weights to torchstore
2. run the generator on the GPU cards
3. release generator's GPU memory occupation
4. reload torchtitan's weights and continue training

But I can't see such an application in the demos

Suggest a potential alternative/fix

No response

[yuchen-zhu-zyc]

I don't think the collocation is supported in the current torchforge I faced the same issue and had to implement my own version.

[shendiaomo]

I don't think the collocation is supported in the current torchforge I faced the same issue and had to implement my own version.

Great, I'm trying the same thing. Can you give a reference implementation as a PR or on your own fork?