[refactor] dynamically import TrainSpec

Author: tianyu-l

torchtitan/__init__.py

@@ -6,8 +6,3 @@
 
 # Import to register quantization modules.
 import torchtitan.components.quantization  # noqa: F401
-
-# Import the built-in models here so that the corresponding register_model_spec()
-# will be called.
-import torchtitan.experiments  # noqa: F401
-import torchtitan.models  # noqa: F401

torchtitan/experiments/__init__.py

@@ -4,7 +4,4 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import torchtitan.experiments.llama4  # noqa: F401
-import torchtitan.experiments.qwen3
-import torchtitan.experiments.simple_fsdp  # noqa: F401
-import torchtitan.experiments.vlm  # noqa: F401
+_supported_experiments = ["flux", "llama4", "qwen3", "simple_fsdp", "vlm"]

torchtitan/experiments/flux/__init__.py

@@ -9,7 +9,7 @@
 
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers
-from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
+from torchtitan.protocols.train_spec import TrainSpec
 
 from .dataset.flux_dataset import build_flux_dataloader
 from .infra.parallelize import parallelize_flux
@@ -107,8 +107,8 @@
 }
 
 
-register_train_spec(
-    TrainSpec(
+def get_train_spec() -> TrainSpec:
+    return TrainSpec(
         name="flux",
         model_cls=FluxModel,
         model_args=flux_configs,
@@ -122,4 +122,3 @@
         build_validator_fn=build_flux_validator,
         state_dict_adapter=FluxStateDictAdapter,
     )
-)

torchtitan/experiments/forge/example_train.py

@@ -297,8 +297,11 @@ def train(self):
                     break
 
                 # Run validation if validator is available
-                if self.job_config.enable and self.validator.should_validate(self.step):
-                    self.validator.validate(self.model_parts)
+                if (
+                    self.job_config.validation.enable
+                    and self.validator.should_validate(self.step)
+                ):
+                    self.validator.validate(self.model_parts, self.step)
 
                 self.checkpointer.save(
                     self.step, last_step=(self.step == job_config.training.steps)

torchtitan/experiments/forge/train_spec.py

@@ -5,12 +5,11 @@
 # LICENSE file in the root directory of this source tree.
 
 from dataclasses import dataclass
+from importlib import import_module
+from typing import Mapping
 
-# Import torchtitan.models to ensure all train specs are registered
-import torchtitan.models  # noqa: F401
 from torchtitan.protocols import BaseModelArgs, BaseStateDictAdapter, ModelProtocol
 from torchtitan.protocols.train_spec import (
-    _train_specs,
     LossFunctionBuilder,
     LRSchedulersBuilder,
     OptimizersBuilder,
@@ -24,7 +23,7 @@
 class ForgeTrainSpec:
     name: str
     model_cls: type[ModelProtocol]
-    model_args: dict[str, BaseModelArgs]
+    model_args: Mapping[str, BaseModelArgs]
     parallelize_fn: ParallelizeFunction
     pipelining_fn: PipeliningFunction | None
     build_optimizers_fn: OptimizersBuilder
@@ -33,24 +32,7 @@ class ForgeTrainSpec:
     state_dict_adapter: type[BaseStateDictAdapter] | None = None
 
 
-# Copy and transform train specs from torchtitan.protocols.train_spec._train_specs
-# This happens during import after all models have been registered
-_forge_train_specs = {}
-
-
-def register_train_spec(train_spec: ForgeTrainSpec) -> None:
-    global _forge_train_specs
-    if train_spec.name in _forge_train_specs:
-        raise ValueError(f"Model {train_spec.name} is already registered.")
-
-    _forge_train_specs[train_spec.name] = train_spec
-
-
-def get_train_spec(name: str) -> ForgeTrainSpec:
-    global _forge_train_specs
-    if name not in _forge_train_specs:
-        raise ValueError(f"Model {name} is not registered.")
-    return _forge_train_specs[name]
+_extra_train_specs = {}
 
 
 def _transform_train_spec(original_spec: TrainSpec):
@@ -69,6 +51,29 @@ def _transform_train_spec(original_spec: TrainSpec):
     )
 
 
-# Populate _forge_train_specs with transformed specs
-for name, spec in _train_specs.items():
-    register_train_spec(_transform_train_spec(spec))
+def register_train_spec(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.")
+
+    # user can define a ForgeTrainSpec from outside of torchtitan
+    _extra_train_specs[train_spec.name] = train_spec
+
+
+def get_train_spec(name: str) -> ForgeTrainSpec:
+    # user-defined ForgeTrainSpec has higher priority
+    global _extra_train_specs
+    if name in _extra_train_specs:
+        return _extra_train_specs[name]
+
+    from torchtitan.experiments import _supported_experiments
+    from torchtitan.models import _supported_models
+
+    if name in _supported_models:
+        module = import_module(f"torchtitan.models.{name}")
+        return _transform_train_spec(module.get_train_spec())
+    elif name in _supported_experiments:
+        module = import_module(f"torchtitan.experiments.{name}")
+        return _transform_train_spec(module.get_train_spec())
+
+    raise ValueError(f"ForgeTrainSpec {name} is not registered.")

torchtitan/experiments/llama4/__init__.py

@@ -11,7 +11,7 @@
 from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.models.llama3 import pipeline_llama
 from torchtitan.models.moe import MoEArgs
-from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
+from torchtitan.protocols.train_spec import TrainSpec
 
 from .infra.parallelize import parallelize_llama
 from .model.args import TransformerModelArgs
@@ -97,8 +97,8 @@
 }
 
 
-register_train_spec(
-    TrainSpec(
+def get_train_spec() -> TrainSpec:
+    return TrainSpec(
         name="llama4",
         model_cls=Transformer,
         model_args=llama4_configs,
@@ -111,4 +111,3 @@
         build_loss_fn=build_cross_entropy_loss,
         state_dict_adapter=Llama4StateDictAdapter,
     )
-)

torchtitan/experiments/simple_fsdp/__init__.py

@@ -12,13 +12,14 @@
 from torchtitan.components.tokenizer import build_hf_tokenizer
 from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.models.llama3 import llama3_configs, pipeline_llama
-from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
+from torchtitan.protocols.train_spec import TrainSpec
 
 from .model import SimpleFSDPTransformer
 from .parallelize import parallelize_llama
 
-register_train_spec(
-    TrainSpec(
+
+def get_train_spec() -> TrainSpec:
+    return TrainSpec(
         name="llama3_simple_fsdp",
         model_cls=SimpleFSDPTransformer,
         model_args=llama3_configs,
@@ -30,4 +31,3 @@
         build_tokenizer_fn=build_hf_tokenizer,
         build_loss_fn=build_cross_entropy_loss,
     )
-)

torchtitan/experiments/vlm/__init__.py

@@ -12,7 +12,7 @@
 from torchtitan.components.tokenizer import build_hf_tokenizer
 from torchtitan.components.validate import build_validator
 from torchtitan.models.llama3 import llama3_configs
-from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
+from torchtitan.protocols.train_spec import TrainSpec
 
 from .datasets.mm_datasets import build_mm_dataloader
 from .infra.parallelize import parallelize_vlm
@@ -40,8 +40,8 @@
 }
 
 
-register_train_spec(
-    TrainSpec(
+def get_train_spec() -> TrainSpec:
+    return TrainSpec(
         name="llama3-siglip2",
         model_cls=Llama3Siglip2Transformer,
         model_args=llama3_siglip2_configs,
@@ -54,4 +54,3 @@
         build_loss_fn=build_cross_entropy_loss,
         build_validator_fn=build_validator,
     )
-)

torchtitan/experiments/vlm/requirements.txt

@@ -1 +1,2 @@
 einops
+pillow

torchtitan/models/README.md

@@ -39,14 +39,16 @@ The folder should be organized as follows
   - Include other util files if necessary.
 - `__init__.py`
   - A dictionary of the actual model configurations, of the type `[str: ModelArgs]`.
-  - Call `register_train_spec` to specify a [`TrainSpec`](/torchtitan/protocols/train_spec.py), consisting a tuple of
+  - Define `get_train_spec` to return a [`TrainSpec`](/torchtitan/protocols/train_spec.py), consisting a tuple of
     - model name, model class, model args
+      - Model name should be the same as the folder name, which should be added to `torchtitan/models/__init__.py` or ``torchtitan/experiments/__init__.py``.
     - parallelizing function, pipelining function
     - builder functions for optimizer, lr scheduler, data loader, tokenizer, and loss function
       - More often than not, existing components can be reused.
       - Adding new datasets requires the `torchtitan` team’s review and legal approval.
       - Try to have minimal dependency on external libraries, if any.
     - state dict adapter
+  - If developing outside of torchtitan, one can call `register_train_spec` to register a `TrainSpec` so that `train.py` can be reused.
   - Read [more](/docs/extension.md#trainspec) on `TrainSpec`.
 - `README.md`
   - Include [instructions](/README.md#downloading-a-tokenizer) to download tokenizers / encoders.