refactor to make model folder structure consistent (#1298)

This PR does the following refactoring:
1. remove `from_model_args` in `ModelProtocal`; add `__init__` and
`init_weights` to `ModelProtocol`
2. make the structure consistent for each model folder:
- model
  - args.py
  - model.py
- infra
  - parallelize.py
  - pipeline.py (optional)

Will publish guidelines on adding new models soon.

Author: tianyu-l

README.md

@@ -77,9 +77,9 @@ We report [performance](benchmarks/llama3_h100_202412_torchtitan.md) on up to 51
 
 You may want to see how the model is defined or how parallelism techniques are applied. For a guided tour, see these files first:
 * [torchtitan/train.py](torchtitan/train.py) - the main training loop and high-level setup code
-* [torchtitan/models/llama3/model.py](torchtitan/models/llama3/model.py) - the Llama 3.1 model definition
-* [torchtitan/models/llama3/parallelize_llama.py](torchtitan/models/llama3/parallelize_llama.py) - helpers for applying Data Parallel, Tensor Parallel, activation checkpointing, and `torch.compile` to the model
-* [torchtitan/models/llama3/pipeline_llama.py](torchtitan/models/llama3/pipeline_llama.py) - helpers for applying Pipeline Parallel to the model
+* [torchtitan/models/llama3/model/model.py](torchtitan/models/llama3/model/model.py) - the Llama 3.1 model definition
+* [torchtitan/models/llama3/infra/parallelize.py](torchtitan/models/llama3/infra/parallelize.py) - helpers for applying Data Parallel, Tensor Parallel, activation checkpointing, and `torch.compile` to the model
+* [torchtitan/models/llama3/infra/pipeline.py](torchtitan/models/llama3/infra/pipeline.py) - helpers for applying Pipeline Parallel to the model
 * [torchtitan/components/checkpoint.py](torchtitan/components/checkpoint.py) - utils for saving/loading distributed checkpoints
 * [torchtitan/components/quantization/float8.py](torchtitan/components/quantization/float8.py) - utils for applying Float8 techniques
 

scripts/estimate/estimation.py

@@ -99,7 +99,7 @@ def estimate_memory(job_config: JobConfig):
             f"Building {train_spec.name} {job_config.model.flavor} with {model_args}"
         )
         with torch.device("meta"):
-            model = model_cls.from_model_args(model_args)
+            model = model_cls(model_args)
 
         # Build the collection of model converters. No-op if `model.converters` empty
         model_converters = build_model_converters(job_config, parallel_dims)

scripts/generate/test_generate.py

@@ -113,7 +113,7 @@ def test_generate(
     init_device = "meta" if world_size > 1 else device
     with torch.device(init_device):
         logger.info(f"Init model on init_device: {init_device}")
-        model = model_cls.from_model_args(model_args)
+        model = model_cls(model_args)
 
     world_mesh = None
     # Init distributed env

tests/unit_tests/test_train_spec.py

@@ -26,10 +26,16 @@
 )
 
 
-class FakeModel(ModelProtocol):
-    @classmethod
-    def from_model_args(cls, args: BaseModelArgs) -> nn.Module:
-        return nn.Linear(8, 8)
+class FakeModel(nn.Module, ModelProtocol):
+    def __init__(self, model_args: BaseModelArgs) -> None:
+        super().__init__()
+        self.linear = nn.Linear(8, 8)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+    def init_weights(self, buffer_device: torch.device | None = None) -> None:
+        nn.init.normal_(self.linear.weight, mean=0.0, std=0.02)
 
 
 def fake_build_optimizers(
@@ -117,7 +123,7 @@ def my_build_optimizer_fn(
 
         apply_to_train_specs(register_optimizer_hook_to_spec)
 
-        model = new_spec.cls.from_model_args(BaseModelArgs())
+        model = new_spec.cls(BaseModelArgs())
         model_parts = [model]
         optimizers = new_spec.build_optimizers_fn(model_parts, JobConfig())
         assert optimizers.optimizers[0].__class__.__name__ == "Adam"

torchtitan/experiments/flux/__init__.py

@@ -9,13 +9,14 @@
 
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers
-from torchtitan.experiments.flux.dataset.flux_dataset import build_flux_dataloader
-from torchtitan.experiments.flux.loss import build_mse_loss
-from torchtitan.experiments.flux.model.autoencoder import AutoEncoderParams
-from torchtitan.experiments.flux.parallelize_flux import parallelize_flux
 from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
 
-from .model.model import FluxModel, FluxModelArgs
+from .dataset.flux_dataset import build_flux_dataloader
+from .infra.parallelize import parallelize_flux
+from .loss import build_mse_loss
+from .model.args import FluxModelArgs
+from .model.autoencoder import AutoEncoderParams
+from .model.model import FluxModel
 
 __all__ = [
     "FluxModelArgs",

torchtitan/experiments/flux/parallelize_flux.py renamed to torchtitan/experiments/flux/infra/parallelize.py

@@ -0,0 +1,37 @@
+# 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.
+
+from dataclasses import dataclass, field
+
+from torch import nn
+
+from torchtitan.experiments.flux.model.autoencoder import AutoEncoderParams
+
+from torchtitan.protocols.train_spec import BaseModelArgs
+from torchtitan.tools.logging import logger
+
+
+@dataclass
+class FluxModelArgs(BaseModelArgs):
+    in_channels: int = 64
+    out_channels: int = 64
+    vec_in_dim: int = 768
+    context_in_dim: int = 512
+    hidden_size: int = 3072
+    mlp_ratio: float = 4.0
+    num_heads: int = 24
+    depth: int = 19
+    depth_single_blocks: int = 38
+    axes_dim: tuple = (16, 56, 56)
+    theta: int = 10_000
+    qkv_bias: bool = True
+    autoencoder_params: AutoEncoderParams = field(default_factory=AutoEncoderParams)
+
+    def get_nparams_and_flops(self, model: nn.Module, seq_len: int) -> tuple[int, int]:
+        # TODO(jianiw): Add the number of flops for the autoencoder
+        nparams = sum(p.numel() for p in model.parameters())
+        logger.warning("FLUX model haven't implement get_nparams_and_flops() function")
+        return nparams, 1

torchtitan/experiments/flux/model/args.py

@@ -4,13 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from dataclasses import dataclass, field
-
 import torch
-
 from torch import nn, Tensor
 
-from torchtitan.experiments.flux.model.autoencoder import AutoEncoderParams
 from torchtitan.experiments.flux.model.layers import (
     DoubleStreamBlock,
     EmbedND,
@@ -20,31 +16,9 @@
     timestep_embedding,
 )
 
-from torchtitan.protocols.train_spec import BaseModelArgs, ModelProtocol
-from torchtitan.tools.logging import logger
-
-
-@dataclass
-class FluxModelArgs(BaseModelArgs):
-    in_channels: int = 64
-    out_channels: int = 64
-    vec_in_dim: int = 768
-    context_in_dim: int = 512
-    hidden_size: int = 3072
-    mlp_ratio: float = 4.0
-    num_heads: int = 24
-    depth: int = 19
-    depth_single_blocks: int = 38
-    axes_dim: tuple = (16, 56, 56)
-    theta: int = 10_000
-    qkv_bias: bool = True
-    autoencoder_params: AutoEncoderParams = field(default_factory=AutoEncoderParams)
-
-    def get_nparams_and_flops(self, model: nn.Module, seq_len: int) -> tuple[int, int]:
-        # TODO(jianiw): Add the number of flops for the autoencoder
-        nparams = sum(p.numel() for p in model.parameters())
-        logger.warning("FLUX model haven't implement get_nparams_and_flops() function")
-        return nparams, 1
+from torchtitan.protocols.train_spec import ModelProtocol
+
+from .args import FluxModelArgs
 
 
 class FluxModel(nn.Module, ModelProtocol):
@@ -159,17 +133,3 @@ def forward(
 
         img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
         return img
-
-    @classmethod
-    def from_model_args(cls, model_args: FluxModelArgs) -> "FluxModel":
-        """
-        Initialize a Flux model from a FluxModelArgs object.
-
-        Args:
-            model_args (FluxModelArgs): Model configuration arguments.
-
-        Returns:
-            FluxModel: FluxModel model.
-
-        """
-        return cls(model_args)

torchtitan/experiments/flux/model/model.py

@@ -16,18 +16,18 @@
 
 from torchtitan.components.tokenizer import Tokenizer
 from torchtitan.config_manager import JobConfig
-from torchtitan.experiments.flux.model.autoencoder import AutoEncoder
+from torchtitan.tools.logging import logger
 
-from torchtitan.experiments.flux.model.hf_embedder import FluxEmbedder
-from torchtitan.experiments.flux.model.model import FluxModel
-from torchtitan.experiments.flux.utils import (
+from .model.autoencoder import AutoEncoder
+from .model.hf_embedder import FluxEmbedder
+from .model.model import FluxModel
+from .utils import (
     create_position_encoding_for_latents,
     generate_noise_latent,
     pack_latents,
     preprocess_data,
     unpack_latents,
 )
-from torchtitan.tools.logging import logger
 
 
 # ----------------------------------------

torchtitan/experiments/flux/sampling.py

@@ -37,7 +37,7 @@ def test_generate_image(self):
         classifier_free_guidance_scale = 5.0
 
         # Contracting JobConfig
-        path = "torchtitan.experiments.flux.flux_argparser"
+        path = "torchtitan.experiments.flux.job_config"
         config_manager = ConfigManager()
         config = config_manager.parse_args(
             [