[Schedule][replace] Transfer hooks when replacing modules

[comaniac]

Description

The current .trace did the following two steps:

1. trace the target module to be a GraphModule.
2. replace the original module with the new GraphModule.

However, the traced GraphModule won't include any hooks in the original module, so the .sync primitives specified before tracing, including .trace and .trace_for_pipeline are dropped. This PR fixes this issue by transferring all hooks from the original module to the new one. For example, a hook registered to broadcast inputs will be transferred 2 times during scheduling:

1. After tracing the top module, the hook is transferred from the original module to the traced module.
2. After wrapping pipeline stage to partitioned modules, the (pre-forward) hook is transferred from the original module to the first pipeline stage module (wrapped by Slapo pipeline module).

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

cc @szhengac @chhzh123

[szhengac]

why don't we also include bwd_pre?

[comaniac]

I don't think there's a backward pre-hook in PyTorch?

[szhengac]

just found this: https://github.com/pytorch/pytorch/blob/b90496eef5665bc39828f6c1c522f399bcc62f3f/torch/nn/modules/module.py#L92. and it seems it does not appear in 1.12