Add an option for automatic module splitting in Autopipeline

• Re "Challenge 3: Forward Method Patching"

Have you looked at:

import torch.distributed.pipelining as pipelining

full_model = AutoModelForCausalLM(...)

pipe = pipelining.pipeline(full_model, spec, ...)
my_submod = pipe.get_stage_module(my_pp_rank)

This way my_submod -- a nn.Module -- would have the desired forward function automatically.

• Re "Challenge 2: nn.ModuleList vs nn.ModuleDict: The Indexing Problem":

my_submod created above would have the same FQN hierarchy as the original model as well as original indices, e.g. layers.8-16 instead of layers.0-8, thus avoiding challenge 2 too.

Reference: Option 2: splitting a model automatically

Originally posted by @kwen2501 in https://github.com/NVIDIA-NeMo/Automodel/discussions/589#discussioncomment-14619715

Start by reading the linked PyTorch distributed.pipelining documentation, especially “Option 2: splitting a model automatically,” and compare it with the Autopipeline code related to the referenced Challenge 2 and Challenge 3. Done means an automatic-splitting option produces submodules with the intended forward function and preserves the original FQN hierarchy and indices.