custom job_config

Author: 3outeille

torchtitan/config/job_config.py

@@ -131,12 +131,6 @@ class Model:
     """
 
 
-@dataclass
-class HFTransformers:
-    model: str = ""
-    """HuggingFace model ID (e.g., 'Qwen/Qwen3-4B-Instruct-2507')"""
-
-
 @dataclass
 class Optimizer:
     name: str = "AdamW"
@@ -903,7 +897,6 @@ class JobConfig:
     profiling: Profiling = field(default_factory=Profiling)
     metrics: Metrics = field(default_factory=Metrics)
     model: Model = field(default_factory=Model)
-    hf_transformers: HFTransformers = field(default_factory=HFTransformers)
     optimizer: Optimizer = field(default_factory=Optimizer)
     lr_scheduler: LRScheduler = field(default_factory=LRScheduler)
     training: Training = field(default_factory=Training)

torchtitan/experiments/transformers_backend/README.md

@@ -17,7 +17,8 @@ hf_assets_path = "./tests/assets/tokenizer"
 +model = "Qwen/Qwen3-4B-Instruct-2507"
 ...
 ```
-- Train: `LOG_RANK=7 CONFIG_FILE=<YOUR_PATHQ/torchtitan/experiments/transformers_backend/configs/qwen3_fsdp2_tp2_pp2.toml ./run_train.sh`
+- Train: `LOG_RANK=7 CONFIG_FILE=<YOUR_PATHQ/torchtitan/experiments/transformers_backend/configs/qwen3_fsdp2_tp2_pp2.toml ./run_train.sh --job.custom_config_module=torchtitan.experiments.transformers_backend.job_config --compile.enable`
+    - Make sure you have created the tokenizers beforehand 
 <img width="1334" height="453" alt="image" src="https://github.com/user-attachments/assets/da459448-027b-4af9-8176-6a3e433a272c" />
 
 ## Supported Features

torchtitan/experiments/transformers_backend/infra/parallelize_hf_transformers.py

@@ -24,7 +24,8 @@
     RowwiseParallel,
     SequenceParallel,
 )
-from torchtitan.config import JobConfig, TORCH_DTYPE_MAP
+from torchtitan.experiments.transformers_backend.job_config import JobConfig
+from torchtitan.config import TORCH_DTYPE_MAP
 from torchtitan.config.job_config import ActivationCheckpoint as ACConfig
 from torchtitan.distributed import NoParallel, ParallelDims
 

torchtitan/experiments/transformers_backend/infra/pipeline_hf.py

@@ -19,7 +19,7 @@
 )
 
 from torchtitan.components.loss import LossFunction
-from torchtitan.config import JobConfig
+from torchtitan.experiments.transformers_backend.job_config import JobConfig
 from torchtitan.distributed import ParallelDims
 from torchtitan.distributed.pipeline_parallel import build_pipeline_schedule
 from torchtitan.protocols.train_spec import BaseModelArgs, ParallelizeFunction

torchtitan/experiments/transformers_backend/job_config.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass, field
+
+@dataclass
+class HFTransformers:
+    model: str = ""
+    """HuggingFace model ID (e.g., 'Qwen/Qwen3-4B-Instruct-2507')"""
+
+@dataclass
+class JobConfig:
+    hf_transformers: HFTransformers = field(default_factory=HFTransformers)

torchtitan/experiments/transformers_backend/model/args.py

@@ -7,7 +7,7 @@
 from dataclasses import dataclass
 
 from torch import nn
-from torchtitan.config import JobConfig
+from torchtitan.experiments.transformers_backend.job_config import JobConfig
 from torchtitan.models.utils import get_dense_model_nparams_and_flops
 from torchtitan.protocols import BaseModelArgs
 from transformers import AutoConfig