from_pretrained distributed refactor (FSDP2 + TP)

[3outeille]

• Introduce DistributedConfig
  • DistributedConfig(tp_size=2, fsdp_size=2) # plans default to "auto" replaces passing separate tp_plan, tp_size, fsdp_plan kwargs. Sizes auto-fill (specify one, the other defaults to 1). Plans default to "auto" when a size is given.
• init_device_mesh + distribute_model
  • init_device_mesh(distributed_config) → builds a 2D DeviceMesh("fsdp", "tp"), inits process group if needed
  • distribute_model(model, distributed_config, device_mesh) → attaches TP hooks before weight loading
  • apply_fsdp2(model, mesh, plan) → wraps with FSDP2 after weight loading
• Made sure accelerate and transformers.distributed dont overlap. Separate codepath

	distributed_config path	accelerate path
Device setup	init_device_mesh()	check_and_set_device_map()
Pre-loading	distribute_model() (TP hooks)	distribute_model() + _get_device_map()
Post-loading	apply_fsdp2()	accelerate_dispatch()

Note: apply_fsdp2 will be moved to pre-loading soon as we dont want to to apply it full weights.

We can now train like this:

# torchrun --nproc_per_node=4 train_fsdp_tp.py

import os
import torch
import torch.distributed.checkpoint as dcp
from torch.utils.data import DataLoader, DistributedSampler
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed import DistributedConfig

if __name__ == "__main__":

    model_name = "meta-llama/Llama-3.2-1B"
    num_steps, lr = 50, 3e-4
    save_dir = "./checkpoints"

    torch.distributed.init_process_group(backend="nccl")
    rank, local_rank = int(os.environ["RANK"]), int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
    tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token

    dataset = load_dataset("roneneldan/TinyStories", split="train[:1000]")
    dataset = dataset.map(lambda x: tokenizer(x["text"], truncation=True, padding="max_length", max_length=512), batched=True, remove_columns=dataset.column_names)
    dataset.set_format("torch")
    dataloader = DataLoader(dataset, batch_size=4, sampler=DistributedSampler(dataset))

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        distributed_config=DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto"),
        torch_dtype=torch.bfloat16,
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    model.train()
    for step, batch in enumerate(dataloader):
        if step >= num_steps:
            break
        input_ids = batch["input_ids"].to(f"cuda:{local_rank}")
        labels = input_ids.clone()
        labels[labels == tokenizer.pad_token_id] = -100

        loss = model(input_ids, labels=labels).loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if rank == 0 and step % 10 == 0:
            print(f"Step {step:>4d} | Loss: {loss.item():.4f}")

    # Save model (HF format) and optimizer (DCP)
    model.save_pretrained(save_dir)
    dcp.save({"optimizer": optimizer.state_dict()}, checkpoint_id=os.path.join(save_dir, "optimizer"))

    if rank == 0:
        tokenizer.save_pretrained(save_dir)
        print(f"Saved to {save_dir}")

    torch.distributed.destroy_process_group()

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: clap, deit

[github-actions]

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=44996&sha=0b8b77