[RFC] Implement model-specific 4d parallelism

[yzhangcs]

Proposal

• We want to add apply_tp & apply_cp fns for each models as their layer definitions can be varied.

Also see comments in fla-org/flame#4

[xffxff]

Things would be easier after pytorch/torchtitan#814 is merged. This PR introduces ModelSpec to describe a model and how to parallelize a model. Here is an example of a ModelSpec for llama3

 ModelSpec(
        name="llama3",
        cls=Transformer,
        config=llama3_configs,
        tokenizer="tiktoken",
        parallelize_fn=parallelize_llama,
        pipelining_fn=pipeline_llama,
    )

We can have different ModelSpecs for fla models.