Fix FFN depth-scaled initialization target

Author: GodlyDonuts

tests/unit_tests/cpu/test_feed_forward.py

@@ -10,8 +10,14 @@
 import torch
 import torch.nn.functional as F
 
+import torchtitan.models.llama3.flavors as llama3_flavors
+
 from torchtitan.models.common.activation import SiTUGLU
-from torchtitan.models.common.config_utils import fused_gate_up_param_init
+from torchtitan.models.common.config_utils import (
+    fused_gate_up_param_init,
+    make_ffn_config,
+    make_shared_expert_ffn_config,
+)
 from torchtitan.models.common.feed_forward import FeedForward
 from torchtitan.models.common.linear import Linear
 
@@ -59,6 +65,62 @@ def test_feed_forward_uses_one_physical_gate_up_linear():
     torch.testing.assert_close(w13_2HD[1], 5 * torch.ones_like(w13_2HD[1]))
 
 
+def test_make_ffn_config_uses_input_init_for_gate_and_up():
+    config = make_ffn_config(
+        dim=4,
+        hidden_dim=8,
+        w1_param_init={"weight": _fill(1.0)},
+        w2_param_init={"weight": _fill(2.0)},
+    )
+    feed_forward = config.build()
+    feed_forward.init_states()
+
+    w13_2HD = feed_forward.w13.weight
+    torch.testing.assert_close(w13_2HD[0], torch.ones_like(w13_2HD[0]))
+    torch.testing.assert_close(w13_2HD[1], torch.ones_like(w13_2HD[1]))
+    torch.testing.assert_close(
+        feed_forward.w2.weight, 2 * torch.ones_like(feed_forward.w2.weight)
+    )
+
+
+def test_make_shared_expert_ffn_config_uses_input_init_for_gate_and_up():
+    config = make_shared_expert_ffn_config(
+        dim=4,
+        hidden_dim=8,
+        w1_param_init={"weight": _fill(1.0)},
+        w2_param_init={"weight": _fill(2.0)},
+    )
+    feed_forward = config.build()
+    feed_forward.init_states()
+
+    w13_2HD = feed_forward.w13.weight
+    torch.testing.assert_close(w13_2HD[0], torch.ones_like(w13_2HD[0]))
+    torch.testing.assert_close(w13_2HD[1], torch.ones_like(w13_2HD[1]))
+    torch.testing.assert_close(
+        feed_forward.w2.weight, 2 * torch.ones_like(feed_forward.w2.weight)
+    )
+
+
+def test_llama3_depth_scales_only_ffn_output(monkeypatch):
+    linear_init = {"weight": _fill(1.0), "bias": _fill(0.0)}
+    depth_init = {"weight": _fill(2.0), "bias": _fill(0.0)}
+    monkeypatch.setattr(llama3_flavors, "_LINEAR_INIT", linear_init)
+    monkeypatch.setattr(
+        llama3_flavors, "_depth_init", lambda _layer_id: depth_init
+    )
+
+    build_config, max_context_length = llama3_flavors.MODEL_FLAVORS["debugmodel"]
+    model_config = build_config(attn_backend="flex", seq_len=max_context_length)
+    feed_forward = model_config.layers[0].feed_forward.build()
+    feed_forward.init_states()
+
+    w13_2HD = feed_forward.w13.weight
+    torch.testing.assert_close(w13_2HD[0], torch.ones_like(w13_2HD[0]))
+    torch.testing.assert_close(w13_2HD[1], torch.ones_like(w13_2HD[1]))
+    torch.testing.assert_close(
+        feed_forward.w2.weight, 2 * torch.ones_like(feed_forward.w2.weight)
+    )
+
 def test_feed_forward_requires_two_w13_projections():
     config = FeedForward.Config(
         w13=Linear.Config(in_features=4, out_features=8),

tests/unit_tests/cpu/test_lora.py

@@ -391,7 +391,7 @@ def test_stacked_lora_adapter_does_not_repeat_base_redistribution():
         dim=4,
         hidden_dim=8,
         w1_param_init=init,
-        w2w3_param_init=init,
+        w2_param_init=init,
     )
     config = LoRATransform(
         rank=2,

tests/unit_tests/cpu/test_moe.py

@@ -428,7 +428,7 @@ def test_shared_expert_uses_runtime_sp_aware_output_projection(self):
             dim=4,
             hidden_dim=8,
             w1_param_init={},
-            w2w3_param_init={},
+            w2_param_init={},
         )
 
         self.assertIs(type(config.w13), ColumnParallelLinear.Config)

tests/unit_tests/cpu/test_tp_kv_heads_validation.py

@@ -93,7 +93,7 @@ def _make_llama3_config(n_heads: int, n_kv_heads: int | None) -> "Llama3Model.Co
                     dim=_DIM,
                     hidden_dim=compute_ffn_hidden_dim(_DIM, multiple_of=256),
                     w1_param_init=_LINEAR_INIT,
-                    w2w3_param_init=_LINEAR_INIT,
+                    w2_param_init=_LINEAR_INIT,
                 ),
             )
         )

tests/unit_tests/gpu/test_async_linear.py

@@ -204,7 +204,7 @@ def test_w13_tp_shards_the_matrix_row_dimension(self):
             dim=DIM,
             hidden_dim=hidden_dim,
             w1_param_init=init,
-            w2w3_param_init=init,
+            w2_param_init=init,
         )
         set_dense_ffn_sharding(
             ffn_config,
@@ -237,7 +237,7 @@ def test_stacked_w13_typechecks(self):
             dim=DIM,
             hidden_dim=128,
             w1_param_init=init,
-            w2w3_param_init=init,
+            w2_param_init=init,
         )
         set_dense_ffn_sharding(
             ffn_config,
@@ -429,7 +429,7 @@ def test_matches_standard_feed_forward(self):
         torch.manual_seed(0)
         standard = (
             make_ffn_config(
-                dim=dim, hidden_dim=hidden, w1_param_init=init, w2w3_param_init=init
+                dim=dim, hidden_dim=hidden, w1_param_init=init, w2_param_init=init
             )
             .build()
             .to(dev)
@@ -438,7 +438,7 @@ def test_matches_standard_feed_forward(self):
             dim=dim,
             hidden_dim=hidden,
             w1_param_init=init,
-            w2w3_param_init=init,
+            w2_param_init=init,
         )
         async_config = AsyncTensorParallelTransform(
             enable_sequence_parallel=True
@@ -534,7 +534,7 @@ def make():
                 dim=dim,
                 hidden_dim=hidden,
                 w1_param_init=init,
-                w2w3_param_init=init,
+                w2_param_init=init,
             )
 
         torch.manual_seed(0)

tests/unit_tests/gpu/test_fused_swiglu.py

@@ -147,7 +147,7 @@ def _dist_gemm_ffn_config(**kwargs):
 
     init = {"weight": torch.nn.init.zeros_}
     return make_ffn_config(
-        dim=_DIM, hidden_dim=_HIDDEN, w1_param_init=init, w2w3_param_init=init, **kwargs
+        dim=_DIM, hidden_dim=_HIDDEN, w1_param_init=init, w2_param_init=init, **kwargs
     )
 
 

tests/unit_tests/gpu/test_tensor_parallel_feed_forward.py

@@ -49,7 +49,7 @@ def test_matches_unsharded_with_and_without_sequence_parallel(self):
                     dim=dim,
                     hidden_dim=hidden_dim,
                     w1_param_init=init,
-                    w2w3_param_init=init,
+                    w2_param_init=init,
                 )
                 reference = copy.deepcopy(base_config).build().to(device)
                 parallel_config = copy.deepcopy(base_config)
@@ -142,7 +142,7 @@ def test_shared_expert_without_sp_reduces_at_parent_boundary(self):
             dim=dim,
             hidden_dim=hidden_dim,
             w1_param_init=init,
-            w2w3_param_init=init,
+            w2_param_init=init,
         )
         reference = copy.deepcopy(base_config).build().to(device)
         parallel_config = copy.deepcopy(base_config)
@@ -220,7 +220,7 @@ def test_shared_expert_with_sp_reduce_scatters_at_w2_boundary(self):
             dim=dim,
             hidden_dim=hidden_dim,
             w1_param_init=init,
-            w2w3_param_init=init,
+            w2_param_init=init,
         )
         reference = copy.deepcopy(base_config).build().to(device)
         parallel_config = copy.deepcopy(base_config)

torchtitan/experiments/transformers_modeling_backend/moe_replacement.py

@@ -530,7 +530,7 @@ def _build_moe_config(params: dict, config) -> MoE.Config:
             dim=shared_info["dim"],
             hidden_dim=shared_info["hidden_dim"],
             w1_param_init=_LINEAR_INIT,
-            w2w3_param_init=_LINEAR_INIT,
+            w2_param_init=_LINEAR_INIT,
         )
         if shared_info["has_sigmoid_gate"]:
             # Import only for the Qwen3.5 topology so unrelated HF models do

torchtitan/models/common/config_utils.py

@@ -265,20 +265,24 @@ def make_ffn_config(
     dim: int,
     hidden_dim: int,
     w1_param_init: dict[str, Callable],
-    w2w3_param_init: dict[str, Callable],
+    w2_param_init: dict[str, Callable],
 ) -> FeedForward.Config:
-    """Build a fully-specified FeedForward.Config."""
+    """Build a fully-specified FeedForward.Config.
+
+    ``w1`` and ``w3`` are the gate/up projections and share
+    ``w1_param_init``; ``w2`` is the residual output projection.
+    """
     return FeedForward.Config(
         w13=ColumnParallelLinear.Config(
             in_features=dim,
             out_features=hidden_dim,
             num_linears=2,
-            param_init=fused_gate_up_param_init(w1_param_init, w2w3_param_init),
+            param_init=fused_gate_up_param_init(w1_param_init, w1_param_init),
         ),
         w2=RowParallelLinear.Config(
             in_features=hidden_dim,
             out_features=dim,
-            param_init=w2w3_param_init,
+            param_init=w2_param_init,
         ),
     )
 
@@ -288,24 +292,27 @@ def make_shared_expert_ffn_config(
     dim: int,
     hidden_dim: int,
     w1_param_init: dict[str, Callable],
-    w2w3_param_init: dict[str, Callable],
+    w2_param_init: dict[str, Callable],
 ) -> FeedForward.Config:
-    """Build a shared FFN whose output reduction is selected at runtime."""
+    """Build a shared FFN whose output reduction is selected at runtime.
+
+    ``w1`` and ``w3`` are the gate/up projections and share
+    ``w1_param_init``; ``w2`` is the residual output projection.
+    """
     return FeedForward.Config(
         w13=ColumnParallelLinear.Config(
             in_features=dim,
             out_features=hidden_dim,
             num_linears=2,
-            param_init=fused_gate_up_param_init(w1_param_init, w2w3_param_init),
+            param_init=fused_gate_up_param_init(w1_param_init, w1_param_init),
         ),
         w2=SharedExpertRowParallelLinear.Config(
             in_features=hidden_dim,
             out_features=dim,
-            param_init=w2w3_param_init,
+            param_init=w2_param_init,
         ),
     )
 
-
 def make_moe_config(
     *,
     num_experts: int = 8,

torchtitan/models/deepseek_v3/flavors.py

@@ -84,7 +84,7 @@ def _depth_experts_init(layer_id: int) -> dict[str, Callable]:
     return {
         "w1_EFD": partial(nn.init.trunc_normal_, std=0.02),
         "w2_EDF": partial(nn.init.trunc_normal_, std=depth_scaled_std(0.02, layer_id)),
-        "w3_EFD": partial(nn.init.trunc_normal_, std=depth_scaled_std(0.02, layer_id)),
+        "w3_EFD": partial(nn.init.trunc_normal_, std=0.02),
     }
 
 
@@ -276,7 +276,7 @@ def build_mla_moe_layers(
                 dim=dim,
                 hidden_dim=dense_hidden_dim,
                 w1_param_init=linear_init,
-                w2w3_param_init=depth_init(layer_id),
+                w2_param_init=depth_init(layer_id),
             )
             moe_cfg = None
         else:
@@ -303,7 +303,7 @@ def build_mla_moe_layers(
                     dim=dim,
                     hidden_dim=moe_hidden_dim * num_shared_experts,
                     w1_param_init=linear_init,
-                    w2w3_param_init=depth_init(layer_id),
+                    w2_param_init=depth_init(layer_id),
                 ),
                 aux_loss_coeff=aux_loss_coeff,
             )