Add Dynamic Model Import and ModelSpec Definition

[fegin]

Stack from ghstack (oldest at bottom):

• Integrate TorchFT #834
• Allow users to customize dataloader #836
• -> Add Dynamic Model Import and ModelSpec Definition #814

What does this PR do?

1. This PR introduces ModelSpec to describe a model and how to parallelize a model.
   • All the models should call register_model_spec().
   • Users can also use --experimental.custom_model_path to dynamically import a model that is not implemented by TorchTitan. The module should also call register_model_spec().
2. This PR also refactors OptimizersContainer and LRSchedulersContainers
   • Fixes an issue that optimizers will accept parameters that requires_grad is False.
   • Improve typing and docstring.
   • Improve the function and class reusability.
   • OptimizersContainer now inherits from torch.optim.Optimizer .
3. This PR also moves parallelize_llama and pipelining_llama to the llama folder.

Why do we need this PR?
This allows users to use TorchTitan with a new model without intrusively change TorchTitan code.

Next steps

1. Dataloader is not included
2. Checkpoint customization is not included yet.

[tianyu-l]

Initial pass looks great. Had some suggestions on restructuring.

[tianyu-l]

LGTM!

[tianyu-l]

This is applying obliviously to all TrainSpec in _train_specs. Now we only have one spec (llama3) and one converter (float8); down the road we may have some converter not working with all specs.

Practically it won't fail since it only defines new functions but won't actually run them. Conceptually it can be a little confusing. But allowing too much complexity is also confusing, so I think it is OK for now.

[tianyu-l]

I can imagine that later there could be cases where in a TrainSpec pipeline_fn is none, so we don't need to implement an optimizer supporting multiple model_parts. If we hit those scenarios, we can consider relax this List[nn.Module].

[fegin]

I don't think we should relax this. Our trainer currently use List[nn.Module] even pipeline degree is 1. This will simplify the code logic.

[tianyu-l]

This is because of our specific way of composing PP and SPMD parallelisms. Later on we may need to revisit such protocols, e.g. how do we do DualPipe from DeepSeek V3 which likely will "fuse" PP and other parallelisms? One way is to put pipeline_fn and parallelize_fn into a bigger customizable concept.

[fegin]

When you say relaxing List[nn.Moulde], what specific type do you have in your mind?

[tianyu-l]

suppose a "community" TrainSpec has a customized pipeline_fn which don't produce multiple model_parts, e.g. something like 1f1b / gpipe, then their corresponding build_optimizer_fn only needs input of type nn.Module, not List[nn.Moulde].

[fegin]

ye, we previously utilized Optional[nn.Module, List[nn.Module]]. But since we changed the trainer to always use model_parts[0] when there is only one nn.Module, we now use List[nn.Module]. We can change it back if that feel more friendly. That will also cause the complexity of the typing as users may have to do cast.