[export] Add customization API for unflattener to control restored module types.

[zhxchen17]

🚀 The feature, motivation and pitch

Today when we call unflattener, the restored module will look like the following:

(UnflattenedModule):
   (InterpreterModule)
      (InterpreterModule)
      (InterpreterModule)
      (InterpreterModule)

This prevents some eager transformations from functioning e.g. pipeline + tensor parallelism or quantization. Ideally it should be possible for users to restore the original submodule types like the following:

(UnflattenedModule):
   (ModuleList)
      (Linear)
      (Linear)
      (Linear)

We had an internal discussion recorded in this doc: https://docs.google.com/document/d/1CAKtSiPWMyjOwaEbTmoT1yM2vF2ZucxxbVeFYhhdy8A/edit

This issue is filed for tracking purpose.

Alternatives

No response

Additional context

No response

cc @avikchaudhuri @gmagogsfm @tugsbayasgalan @angelayi @suo @ydwu4