[P0] To Support DDP/FSDP (all stage)

[frankaging]

Suggestion / Feature Request

Descriptions:

Currently, one of our downstream library pyreft enables DDP by doing work-around, since pyvene intervenable stores interventions in a python-native dict which cannot be picked up by torch DP wrappers as parameters. This is one of the major blocker of DP support.

This ticket will involve a set of changes in pyvene backbone to support different kinds of DP. This will unlock pyvene support of LLMs larger than 70B for training and inference w/ interventions.