refactor TrainSpec to remove the name field (#1850)

Author: tianyu-l

docs/extension.md

@@ -14,7 +14,7 @@ The extension points and protocols mentioned in this note are subject to change.
 The coarse level abstraction tries to hit a balance between flexible component swapping and a straightforward train script ([train.py](../torchtitan/train.py)).
 Note that among all training components, currently [`CheckpointManager`](../torchtitan/components/checkpoint.py) and [`FTManager`](../torchtitan/components/ft.py) are not configurable since we do not expect them to be customized, but we are open to requests.
 
-To register a `TrainSpec`, please follow the example of [Llama 3.1](../torchtitan/models/llama3/__init__.py) to `register_train_spec`. Please make sure the registration code is called before training initialization. In torchtitan, it is performed during  [module import](../torchtitan/__init__.py).
+To register a `TrainSpec`, please use the `register_train_spec` API, and make sure registration happens before `get_train_spec` is called during training initialization. In torchtitan, `get_train_spec` will dynamically look for models in `torchtitan/models` or `torchtitan/experiments`.
 
 
 ### `ModelConverter`

scripts/estimate/estimation.py

@@ -95,7 +95,7 @@ def estimate_memory(job_config: JobConfig):
         else contextlib.nullcontext()
     ):
         logger.info(
-            f"Building {train_spec.name} {job_config.model.flavor} with {model_args}"
+            f"Building {job_config.model.name} {job_config.model.flavor} with {model_args}"
         )
         with torch.device("meta"):
             model = train_spec.model_cls(model_args)

tests/unit_tests/test_train_spec.py

@@ -76,7 +76,6 @@ class TestTrainSpec:
     def test_register_train_spec(self):
         fake_config = {"fake": BaseModelArgs()}
         spec = TrainSpec(
-            name="fake",
             model_cls=FakeModel,
             model_args=fake_config,
             parallelize_fn=parallelize_llama,
@@ -87,7 +86,7 @@ def test_register_train_spec(self):
             build_tokenizer_fn=build_hf_tokenizer,
             build_loss_fn=build_cross_entropy_loss,
         )
-        register_train_spec(spec)
+        register_train_spec("fake", spec)
         new_spec = get_train_spec("fake")
         assert new_spec == spec
 
@@ -98,7 +97,6 @@ def test_optim_hook(self):
         fake_config = {"fake": BaseModelArgs()}
 
         spec = TrainSpec(
-            name="fake2",
             model_cls=FakeModel,
             model_args=fake_config,
             parallelize_fn=parallelize_llama,
@@ -109,7 +107,7 @@ def test_optim_hook(self):
             build_tokenizer_fn=build_hf_tokenizer,
             build_loss_fn=build_cross_entropy_loss,
         )
-        register_train_spec(spec)
+        register_train_spec("fake2", spec)
         new_spec = get_train_spec("fake2")
 
         model = new_spec.model_cls(BaseModelArgs())

torchtitan/experiments/deepseek_v3/__init__.py

@@ -40,8 +40,8 @@
 
 
 register_train_spec(
+    "deepseek3",
     TrainSpec(
-        name="deepseek3",
         model_cls=DeepseekForCausalLM,
         model_args=deepseek_configs,
         parallelize_fn=parallelize_deepseek,
@@ -51,5 +51,5 @@
         build_dataloader_fn=build_hf_dataloader,
         build_tokenizer_fn=get_hf_tokenizer,
         build_loss_fn=build_cross_entropy_loss,
-    )
+    ),
 )

torchtitan/experiments/flux/__init__.py

@@ -109,7 +109,6 @@
 
 def get_train_spec() -> TrainSpec:
     return TrainSpec(
-        name="flux",
         model_cls=FluxModel,
         model_args=flux_configs,
         parallelize_fn=parallelize_flux,

torchtitan/experiments/forge/engine.py

@@ -167,7 +167,7 @@ def __init__(self, job_config: ForgeJobConfig):
         if parallel_dims.pp_enabled:
             if not self.train_spec.pipelining_fn:
                 raise RuntimeError(
-                    f"Pipeline Parallel is enabled but {self.train_spec.name} "
+                    f"Pipeline Parallel is enabled but {job_config.model.name} "
                     f"does not support pipelining"
                 )
 

torchtitan/experiments/forge/example_train.py

@@ -66,7 +66,7 @@ def __init__(self, job_config: JobConfig):
 
         model_args = self.model_args
         logger.info(
-            f"Built {self.train_spec.name} {job_config.model.flavor} with {model_args}"
+            f"Built {job_config.model.name} {job_config.model.flavor} with {model_args}"
         )
 
         # metrics logging
@@ -78,7 +78,7 @@ def __init__(self, job_config: JobConfig):
         self.metrics_processor.num_flops_per_token = self.num_flops_per_token
 
         logger.info(
-            f"{color.blue}Model {self.train_spec.name} {job_config.model.flavor} "
+            f"{color.blue}Model {job_config.model.name} {job_config.model.flavor} "
             f"{color.red}size: {self.model_param_count:,} total parameters{color.reset}"
         )
 

torchtitan/experiments/forge/train_spec.py

@@ -21,7 +21,6 @@
 
 @dataclass
 class ForgeTrainSpec:
-    name: str
     model_cls: type[ModelProtocol]
     model_args: Mapping[str, BaseModelArgs]
     parallelize_fn: ParallelizeFunction
@@ -39,7 +38,6 @@ def _transform_train_spec(original_spec: TrainSpec):
     """Transform the original train spec to ForgeTrainSpec format."""
     # Create a new TrainSpec with only the fields we need in forge
     return ForgeTrainSpec(
-        name=original_spec.name,
         model_cls=original_spec.model_cls,
         model_args=original_spec.model_args,
         parallelize_fn=original_spec.parallelize_fn,
@@ -51,13 +49,13 @@ def _transform_train_spec(original_spec: TrainSpec):
     )
 
 
-def register_train_spec(train_spec: ForgeTrainSpec) -> None:
+def register_train_spec(name: str, train_spec: ForgeTrainSpec) -> None:
     global _extra_train_specs
-    if train_spec.name in _extra_train_specs:
-        raise ValueError(f"ForgeTrainSpec {train_spec.name} is already registered.")
+    if name in _extra_train_specs:
+        raise ValueError(f"ForgeTrainSpec {name} is already registered.")
 
     # user can define a ForgeTrainSpec from outside of torchtitan
-    _extra_train_specs[train_spec.name] = train_spec
+    _extra_train_specs[name] = train_spec
 
 
 def get_train_spec(name: str) -> ForgeTrainSpec:

torchtitan/experiments/llama4/__init__.py

@@ -8,6 +8,7 @@
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers_with_moe_load_balancing
 from torchtitan.components.tokenizer import build_hf_tokenizer
+from torchtitan.components.validate import build_validator
 from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.models.llama3 import pipeline_llama
 from torchtitan.models.moe import MoEArgs
@@ -103,7 +104,6 @@
 
 def get_train_spec() -> TrainSpec:
     return TrainSpec(
-        name="llama4",
         model_cls=Transformer,
         model_args=llama4_configs,
         parallelize_fn=parallelize_llama,
@@ -113,5 +113,6 @@ def get_train_spec() -> TrainSpec:
         build_dataloader_fn=build_hf_dataloader,
         build_tokenizer_fn=build_hf_tokenizer,
         build_loss_fn=build_cross_entropy_loss,
+        build_validator_fn=build_validator,
         state_dict_adapter=Llama4StateDictAdapter,
     )

torchtitan/experiments/multimodal/__init__.py

@@ -22,8 +22,8 @@
 }
 
 register_train_spec(
+    "llama4_multimodal",
     TrainSpec(
-        name="llama4_multimodal",
         model_cls=MultimodalDecoder,
         model_args=llama4_mm_configs,
         parallelize_fn=parallelize_llama,
@@ -33,5 +33,5 @@
         build_dataloader_fn=build_mm_dataloader,
         build_tokenizer_fn=build_hf_tokenizer,
         build_loss_fn=build_cross_entropy_loss,
-    )
+    ),
 )