refactor: AutoModel entry point

[akoumpa]

HF

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
)

NeMo AutoModel

from nemo_automodel import NeMoAutoModelForCausalLM
from torch.distributed.device_mesh import init_device_mesh

mesh = init_device_mesh("cuda", mesh_shape=(1,1,1,1,2), mesh_dim_names=("pp","dp_replicate","dp_shard","cp","tp"))

model = NeMoAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    device_mesh=mesh,
    distributed={"tp_size": 2, "cp_size": 1, "pp_size": 1, "dp_size": 2, "backend": "nccl"},
)

This PR extends the Auto API to include the device_mesh and distributed options. The goal is to provide a drop-in class that supports models with distributed processing.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.