Add transformers backend (Dense model only)

torchtitan/experiments/__init__.py

@@ -7,3 +7,4 @@
 import torchtitan.experiments.llama4  # noqa: F401
 import torchtitan.experiments.qwen3
 import torchtitan.experiments.simple_fsdp  # noqa: F401
+import torchtitan.experiments.transformers_backend # noqa: F401

torchtitan/experiments/transformers_backend/__init__.py

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import dataclasses
+
+from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.components.lr_scheduler import build_lr_schedulers
+from torchtitan.components.optimizer import build_optimizers
+from torchtitan.datasets.hf_datasets import build_hf_dataloader
+from torchtitan.components.tokenizer import build_hf_tokenizer
+
+from torchtitan.models.llama3 import pipeline_llama
+from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
+
+from .infra.parallelize_hf_transformers import parallelize_hf_transformers
+from .model.hf_transformers_args import HFTransformerModelArgs
+
+from transformers.models.llama.modeling_llama import LlamaForCausalLM
+
+
+__all__ = [
+    "HFTransformerModelArgs",
+    "LlamaForCausalLM", #TODO(3outeille): later use AutoModelForCausalLM
+    "hf_transformers_configs",
+]
+
+
+flavors = {
+    "debug": HFTransformerModelArgs(
+        dim=1,
+        n_layers=6,
+        n_heads=16,
+        rope_theta=500000,
+    ),
+    "medium": HFTransformerModelArgs(
+        dim=40,
+        n_layers=24,
+        n_heads=32,
+        rope_theta=500000,
+    ),
+    "full": HFTransformerModelArgs(),
+}
+
+hf_train_spec = TrainSpec(
+    name="hf_auto_model",
+    model_cls=LlamaForCausalLM,
+    model_args=flavors,
+    parallelize_fn=parallelize_hf_transformers,
+    pipelining_fn=pipeline_llama,
+    build_optimizers_fn=build_optimizers,
+    build_lr_schedulers_fn=build_lr_schedulers,
+    build_dataloader_fn=build_hf_dataloader,
+    build_tokenizer_fn=build_hf_tokenizer,
+    build_loss_fn=build_cross_entropy_loss,
+)
+
+# Register multiple train_specs under the same name
+register_train_spec(hf_train_spec)
+register_train_spec(dataclasses.replace(hf_train_spec, name="meta-llama/Llama-3.2-3B"))
+register_train_spec(dataclasses.replace(hf_train_spec, name="meta-llama/Llama-3.2-1B"))

torchtitan/experiments/transformers_backend/configs/debug_1_gpu.toml

@@ -0,0 +1,62 @@
+[job]
+dump_folder = "./outputs"
+description = "Llama 3 debug training with FSDP on 2 GPUs"
+print_args = false
+use_for_integration_test = true
+
+[profiling]
+enable_profiling = false
+save_traces_folder = "profile_trace"
+profile_freq = 10
+enable_memory_snapshot = false
+save_memory_snapshot_folder = "memory_snapshot"
+
+[metrics]
+log_freq = 1
+disable_color_printing = false
+enable_tensorboard = false
+save_tb_folder = "tb"
+enable_wandb = false
+
+[model]
+name = "llama3"
+flavor = "debugmodel"
+tokenizer_path = "/fsx/ferdinandmom/ferdinand-hf/huggingface/torchtitan/tests/assets/tokenizer"
+
+[optimizer]
+name = "AdamW"
+lr = 8e-4
+eps = 1e-8
+
+[lr_scheduler]
+warmup_steps = 2
+decay_ratio = 0.8
+decay_type = "linear"
+min_lr_factor = 0.0
+
+[training]
+local_batch_size = 8
+seq_len = 2048
+max_norm = 1.0
+steps = 10
+compile = false
+dataset = "c4_test"
+dataset_path = "/fsx/ferdinandmom/ferdinand-hf/huggingface/torchtitan/tests/assets/c4_test"
+
+[parallelism]
+data_parallel_replicate_degree = 1
+data_parallel_shard_degree = 1
+tensor_parallel_degree = 1
+pipeline_parallel_degree = 1
+context_parallel_degree = 1
+expert_parallel_degree = 1
+
+[checkpoint]
+enable_checkpoint = false
+
+[activation_checkpoint]
+mode = "selective"
+selective_ac_option = '2'
+
+[validation]
+enabled = false 

torchtitan/experiments/transformers_backend/configs/debug_1_gpu_hf.toml

@@ -0,0 +1,62 @@
+[job]
+dump_folder = "./outputs"
+description = "Llama 3 debug training with FSDP on 2 GPUs"
+print_args = false
+use_for_integration_test = true
+
+[profiling]
+enable_profiling = false
+save_traces_folder = "profile_trace"
+profile_freq = 10
+enable_memory_snapshot = false
+save_memory_snapshot_folder = "memory_snapshot"
+
+[metrics]
+log_freq = 1
+disable_color_printing = false
+enable_tensorboard = false
+save_tb_folder = "tb"
+enable_wandb = false
+
+[model]
+name = "meta-llama/Llama-3.2-1B"
+flavor = "medium"
+tokenizer_path = "/fsx/ferdinandmom/ferdinand-hf/huggingface/torchtitan/tests/assets/tokenizer"
+
+[optimizer]
+name = "AdamW"
+lr = 8e-4
+eps = 1e-8
+
+[lr_scheduler]
+warmup_steps = 2
+decay_ratio = 0.8
+decay_type = "linear"
+min_lr_factor = 0.0
+
+[training]
+local_batch_size = 8
+seq_len = 2048
+max_norm = 1.0
+steps = 10
+compile = false
+dataset = "c4_test"
+dataset_path = "/fsx/ferdinandmom/ferdinand-hf/huggingface/torchtitan/tests/assets/c4_test"
+
+[parallelism]
+data_parallel_replicate_degree = 1
+data_parallel_shard_degree = 1
+tensor_parallel_degree = 1
+pipeline_parallel_degree = 1
+context_parallel_degree = 1
+expert_parallel_degree = 1
+
+[checkpoint]
+enable_checkpoint = false
+
+[activation_checkpoint]
+mode = "selective"
+selective_ac_option = '2'
+
+[validation]
+enabled = false 

torchtitan/experiments/transformers_backend/configs/debug_fsdp_2_gpu.toml

@@ -0,0 +1,65 @@
+# FSDP-only configuration for a 2-GPU setup.
+# Model is sharded across GPUs.
+
+[job]
+dump_folder = "./outputs"
+description = "Llama 3 debug training with FSDP on 2 GPUs"
+print_args = false
+use_for_integration_test = true
+
+[profiling]
+enable_profiling = false
+save_traces_folder = "profile_trace"
+profile_freq = 10
+enable_memory_snapshot = false
+save_memory_snapshot_folder = "memory_snapshot"
+
+[metrics]
+log_freq = 1
+disable_color_printing = false
+enable_tensorboard = false
+save_tb_folder = "tb"
+enable_wandb = false
+
+[model]
+name = "llama3"
+flavor = "debugmodel"
+tokenizer_path = "/fsx/ferdinandmom/ferdinand-hf/huggingface/torchtitan/tests/assets/tokenizer"
+
+[optimizer]
+name = "AdamW"
+lr = 8e-4
+eps = 1e-8
+
+[lr_scheduler]
+warmup_steps = 2
+decay_ratio = 0.8
+decay_type = "linear"
+min_lr_factor = 0.0
+
+[training]
+local_batch_size = 8
+seq_len = 2048
+max_norm = 1.0
+steps = 10
+compile = false
+dataset = "c4_test"
+dataset_path = "/fsx/ferdinandmom/ferdinand-hf/huggingface/torchtitan/tests/assets/c4_test"
+
+[parallelism]
+data_parallel_replicate_degree = 1
+data_parallel_shard_degree = 2
+tensor_parallel_degree = 1
+pipeline_parallel_degree = 1
+context_parallel_degree = 1
+expert_parallel_degree = 1
+
+[checkpoint]
+enable_checkpoint = false
+
+[activation_checkpoint]
+mode = "selective"
+selective_ac_option = '2'
+
+[validation]
+enabled = false