deprecate norm.py and norm_type from JobConfig (#1115)

`norm.py` and `norm_type` in `JobConfig` and llama `ModelArgs` were
introduced when `nn.RMSNorm` was not available. Now that we don't have
such need, let's remove them, following #1111.

Author: tianyu-l

scripts/estimate/estimation.py

@@ -34,10 +34,6 @@ def estimate_memory(job_config: JobConfig):
     # Get the world size
     world_size = int(os.environ["WORLD_SIZE"])
 
-    if job_config.model.norm_type == "compiled_rmsnorm":
-        logger.info("Compiled RMSNorm is not supported yet. Switching to RMSNorm.")
-        job_config.model.norm_type = "rmsnorm"
-
     if job_config.training.compile or job_config.parallelism.enable_compiled_autograd:
         logger.info("Compile mode is not supported yet. Switching to eager mode.")
         job_config.training.compile = False
@@ -91,25 +87,19 @@ def estimate_memory(job_config: JobConfig):
 
     # build model (using meta init)
     model_cls = train_spec.cls
-    model_config = train_spec.config[job_config.model.flavor]
-    # set the model configs from training inputs:
-    # 1. norm type to decide which norm layer to use
-    # 2. vocab size from tokenizer
-    # 3. max_seq_len base on inputs
-    model_config.norm_type = job_config.model.norm_type
-    model_config.vocab_size = tokenizer.n_words
-    model_config.max_seq_len = job_config.training.seq_len
+    model_args = train_spec.config[job_config.model.flavor]
+    model_args.update_from_config(job_config, tokenizer)
 
     with (
         FakeTensorMode()
         if not job_config.memory_estimation.disable_fake_mode
         else contextlib.nullcontext()
     ):
         logger.info(
-            f"Building {train_spec.name} {job_config.model.flavor} with {model_config}"
+            f"Building {train_spec.name} {job_config.model.flavor} with {model_args}"
         )
         with torch.device("meta"):
-            model = model_cls.from_model_args(model_config)
+            model = model_cls.from_model_args(model_args)
 
         # Build the collection of model converters. No-op if `model.converters` empty
         model_converters = build_model_converters(job_config, parallel_dims)
@@ -134,19 +124,19 @@ def estimate_memory(job_config: JobConfig):
             lambda *args, **kwargs: model_converters.post_optimizer_hook(model)
         )
 
-        logger.info(f"Vocab size: {model_config.vocab_size}")
+        logger.info(f"Vocab size: {model_args.vocab_size}")
         # Create a dummy batch instead of loading from a dataset
         batch = (
             torch.randint(
                 0,
-                model_config.vocab_size,
-                (job_config.training.batch_size, model_config.max_seq_len),
+                model_args.vocab_size,
+                (job_config.training.batch_size, model_args.max_seq_len),
                 device="cuda",
             ),
             torch.randint(
                 0,
-                model_config.vocab_size,
-                (job_config.training.batch_size, model_config.max_seq_len),
+                model_args.vocab_size,
+                (job_config.training.batch_size, model_args.max_seq_len),
                 device="cuda",
             ),
         )

scripts/generate/test_generate.py

@@ -85,9 +85,9 @@ def test_generate(
     color = utils.Color
 
     # Load configuration from toml file
-    config = JobConfig()
-    config.parse_args([f"--job.config_file={config_path}"])
-    config._validate_config()
+    job_config = JobConfig()
+    job_config.parse_args([f"--job.config_file={config_path}"])
+    job_config._validate_config()
 
     if len(args.prompt) == 0:
         logger.warning(
@@ -100,27 +100,26 @@ def test_generate(
     device_module.set_device(device)
     device_memory_monitor = build_device_memory_monitor()
 
-    train_spec = get_train_spec(config.model.name)
+    train_spec = get_train_spec(job_config.model.name)
 
     logger.info(f"World Size: {world_size}, Local Rank: {local_rank} on {device}")
 
     # Tokenizer setup
-    tokenizer = train_spec.build_tokenizer_fn(config)
-    model_config = train_spec.config[config.model.flavor]
-    model_config.norm_type = config.model.norm_type
-    model_config.max_seq_len = config.training.seq_len
-    model_config.vocab_size = tokenizer.n_words
+    tokenizer = train_spec.build_tokenizer_fn(job_config)
 
     model_cls = train_spec.cls
+    model_args = train_spec.config[job_config.model.flavor]
+    model_args.update_from_config(job_config, tokenizer)
+
     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_config)
+        model = model_cls.from_model_args(model_args)
 
     world_mesh = None
     # Init distributed env
     if world_size > 1:
-        dist_utils.init_distributed(config)
+        dist_utils.init_distributed(job_config)
         parallel_dims = ParallelDims(
             dp_replicate=1,
             dp_shard=-1,

torchtitan/config_manager.py

@@ -186,13 +186,6 @@ def __init__(self):
             default="debugmodel",
             help="Which model config to train",
         )
-        self.parser.add_argument(
-            "--model.norm_type",
-            type=str,
-            default="rmsnorm",
-            choices=["layernorm", "np_layernorm", "rmsnorm"],
-            help="Type of layer normalization to use [layernorm, np_layernorm, rmsnorm]",
-        )
         self.parser.add_argument(
             "--model.tokenizer_path",
             type=str,

torchtitan/experiments/llama4/model/args.py

@@ -32,7 +32,6 @@ class TransformerModelArgs(BaseModelArgs):
     # If `True`, then each transformer block init uses its layer ID, and if
     # `False`, each uses the total number of transformer blocks
     depth_init: bool = True
-    norm_type: str = "rmsnorm"
 
     use_flex_attn: bool = False
     attn_mask_type: str = "causal"
@@ -59,7 +58,6 @@ class TransformerModelArgs(BaseModelArgs):
     use_grouped_mm: bool = True  # grouped mm or for-loop for the experts computation
 
     def update_from_config(self, job_config: JobConfig, tokenizer: Tokenizer) -> None:
-        self.norm_type = job_config.model.norm_type
         self.vocab_size = tokenizer.n_words
         self.max_seq_len = job_config.training.seq_len
         if self.use_grouped_mm and not has_cuda_capability(9, 0):

torchtitan/experiments/llama4/model/model.py

@@ -10,7 +10,6 @@
 from torch import nn
 
 from torchtitan.models.attention import build_attention, init_attention_mask
-from torchtitan.models.norms import build_norm
 from torchtitan.protocols.train_spec import ModelProtocol
 
 from .args import TransformerModelArgs
@@ -311,20 +310,13 @@ def __init__(
                 ffn_dim_multiplier=model_args.ffn_dim_multiplier,
             )
 
-        self.layer_id = layer_id
-        self.num_layers = model_args.n_layers
-
-        self.attention_norm = build_norm(
-            model_args.norm_type, dim=model_args.dim, eps=model_args.norm_eps
-        )
-        self.ffn_norm = build_norm(
-            model_args.norm_type, dim=model_args.dim, eps=model_args.norm_eps
-        )
+        self.attention_norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
+        self.ffn_norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
 
         if model_args.depth_init:
-            self.weight_init_std = 0.02 / (2 * (self.layer_id + 1)) ** 0.5
+            self.weight_init_std = 0.02 / (2 * (layer_id + 1)) ** 0.5
         else:
-            self.weight_init_std = 0.02 / (2 * self.num_layers) ** 0.5
+            self.weight_init_std = 0.02 / (2 * model_args.n_layers) ** 0.5
 
     def forward(
         self,
@@ -399,11 +391,7 @@ def __init__(self, model_args: TransformerModelArgs):
         self.layers = torch.nn.ModuleDict()
         for layer_id in range(model_args.n_layers):
             self.layers[str(layer_id)] = TransformerBlock(layer_id, model_args)
-
-        self.norm = build_norm(
-            model_args.norm_type, dim=model_args.dim, eps=model_args.norm_eps
-        )
-
+        self.norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
         self.output = nn.Linear(model_args.dim, model_args.vocab_size, bias=False)
         self.init_weights()
 

torchtitan/experiments/llama4/train_configs/debug_model.toml

@@ -21,7 +21,6 @@ enable_wandb = false
 [model]
 name = "llama4"
 flavor = "debugmodel"
-norm_type = "rmsnorm"  # layernorm / np_layernorm / rmsnorm
 # test tokenizer.model, for debug purpose only
 tokenizer_path = "./tests/assets/test_tiktoken.model"
 # converters = "float8"

torchtitan/experiments/llama4/train_configs/llama4_17bx128e.toml

@@ -17,7 +17,6 @@ save_tb_folder = "tb"
 [model]
 name = "llama4"
 flavor = "17bx128e"
-norm_type = "rmsnorm"  # layernorm / np_layernorm / rmsnorm
 tokenizer_path = "./assets/tokenizer/tokenizer.model"
 # converters = "float8"
 

torchtitan/experiments/llama4/train_configs/llama4_17bx16e.toml

@@ -17,7 +17,6 @@ save_tb_folder = "tb"
 [model]
 name = "llama4"
 flavor = "17bx16e"
-norm_type = "rmsnorm"  # layernorm / np_layernorm / rmsnorm
 tokenizer_path = "./assets/tokenizer/tokenizer.model"
 # converters = "float8"
 

torchtitan/experiments/multimodal/model.py

@@ -56,7 +56,6 @@ class ModelArgs:
     # If `True`, then each transformer block init uses its layer ID, and if
     # `False`, each uses the total number of transformer blocks
     depth_init: bool = True
-    norm_type: str = "rmsnorm"
 
 
 class Fp32LayerNorm(nn.LayerNorm):

torchtitan/models/llama3/model.py

@@ -17,7 +17,6 @@
 from torchtitan.components.tokenizer import Tokenizer
 from torchtitan.config_manager import JobConfig
 from torchtitan.models.attention import build_attention, init_attention_mask
-from torchtitan.models.norms import build_norm
 from torchtitan.protocols.train_spec import BaseModelArgs, ModelProtocol
 
 
@@ -37,14 +36,12 @@ class TransformerModelArgs(BaseModelArgs):
     # If `True`, then each transformer block init uses its layer ID, and if
     # `False`, each uses the total number of transformer blocks
     depth_init: bool = True
-    norm_type: str = "rmsnorm"
 
     use_flex_attn: bool = False
     attn_mask_type: str = "causal"
     eos_id: int = 0
 
     def update_from_config(self, job_config: JobConfig, tokenizer: Tokenizer) -> None:
-        self.norm_type = job_config.model.norm_type
         self.vocab_size = tokenizer.n_words
         self.max_seq_len = job_config.training.seq_len
 
@@ -341,20 +338,13 @@ def __init__(self, layer_id: int, model_args: TransformerModelArgs):
             multiple_of=model_args.multiple_of,
             ffn_dim_multiplier=model_args.ffn_dim_multiplier,
         )
-        self.layer_id = layer_id
-        self.num_layers = model_args.n_layers
-
-        self.attention_norm = build_norm(
-            model_args.norm_type, dim=model_args.dim, eps=model_args.norm_eps
-        )
-        self.ffn_norm = build_norm(
-            model_args.norm_type, dim=model_args.dim, eps=model_args.norm_eps
-        )
+        self.attention_norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
+        self.ffn_norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
 
         if model_args.depth_init:
-            self.weight_init_std = 0.02 / (2 * (self.layer_id + 1)) ** 0.5
+            self.weight_init_std = 0.02 / (2 * (layer_id + 1)) ** 0.5
         else:
-            self.weight_init_std = 0.02 / (2 * self.num_layers) ** 0.5
+            self.weight_init_std = 0.02 / (2 * model_args.n_layers) ** 0.5
 
     def forward(
         self,
@@ -423,11 +413,7 @@ def __init__(self, model_args: TransformerModelArgs):
         self.layers = torch.nn.ModuleDict()
         for layer_id in range(model_args.n_layers):
             self.layers[str(layer_id)] = TransformerBlock(layer_id, model_args)
-
-        self.norm = build_norm(
-            model_args.norm_type, dim=model_args.dim, eps=model_args.norm_eps
-        )
-
+        self.norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
         self.output = nn.Linear(model_args.dim, model_args.vocab_size, bias=False)
         self.init_weights()