fix(model): initialize wrapped models, bound trunc_normal_, keep all pipeline stages

Three pre-existing bugs, each producing a silently wrong model rather than an
error.

1. Weight initialization was skipped entirely for wrapped models.
   NamedParameterwiseNormalInitialization and Llama3Initializer stripped only
   torch.compile's `_orig_mod.` prefix before matching their parameter-name
   regexes. Activation checkpointing and FSDP1 insert their own segments
   (`_checkpoint_wrapped_module.`, `_fsdp_wrapped_module.`), so any config that
   wrapped the model before initializing it matched *no* per-layer regex and
   silently kept the default initialization. A shared normalize_parameter_name
   now strips all of them.

2. Llama3Initializer injected one out-of-distribution weight into ~0.3% of
   tensors. Its trunc_normal_ calls passed a=-2, b=2; torch treats those as
   absolute bounds, but the intended standard deviations are 0.02 and smaller,
   so the bounds sat at +-100 to +-283 sigma. That is not just a no-op: the erf
   limits of the inverse-transform sampler saturate, its singular edge becomes
   reachable, and the final clamp pins the affected element to exactly the
   bound - one weight of magnitude 2.0 in a tensor whose intended scale is
   0.007. Measured 12/4000 tensors (0.30%); after the fix 0/4000. Bounds are now
   expressed in standard deviations, matching both the convention this same file
   already used for the output projection (3 / sqrt(n_embd), i.e. 3 * std) and
   the llama3 reference implementation.

   This is the root cause of the intermittent TestLlama3LikeInitialization
   failures: the pinned element dominates the sample variance, which is what the
   test's std assertions detect. Measured 6 failures in 40 runs (15%) before,
   0 in 60 after. The test's max/min assertions were pinned to the old absolute
   bound; they now assert 3 * std, so they detect a stray element rather than
   tolerating one.

3. Pipeline stage generation silently dropped modules. get_stages packed split
   points greedily against a fixed per-stage weight cap while looping exactly
   num_virtual_stages times; whatever remained when the packing did not fit was
   never assigned - in practice the output split point, yielding a pipeline with
   no lm_head on any stage. Uniform per-layer weights always happen to fit,
   which is why GPT2 never hit it (verified: 0 of ~1000 GPT2 configurations drop
   anything), but any generator with non-uniform per-layer cost does. Packing
   now uses a partitioner that provably assigns every module exactly once, and
   requesting more stages than split points raises instead of returning empty
   stages.

Adds tests/models/parallelism/test_stages_generator.py and
tests/nn/model_initialization/test_fqn_normalization.py. Five of the stage tests
fail against the old packer, including one asserting lm_head survives; GPT2's
cases pass either way, confirming GPT2 was unaffected.

Note: item 2 changes the numerics of Llama3Initializer. Seeded runs from before
and after will not produce identical weights.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: le1nux

CHANGELOG_DEV.md

@@ -4,6 +4,7 @@
 |------------------|------------|---------------|------------------|------------------------------------------------------------------------------------------------|
 | [#141](#pr-141-towards-stable-modalities-version)  | Bug Fix    |  [#129](https://github.com/Modalities/modalities/issues/129)         | **Yes**              | Towards stable modalities version                                                               |
 | [#154](pr-154-manual-swiglu-implementation)  | Bug Fix    |  [#14](https://github.com/Modalities/modalities/issues/14)         | **Yes**              | Towards stable modalities version                                                               |
+| [#init-and-pipeline-fixes](#pr-initialization-and-pipeline-stage-fixes)  | Bug Fix    |  --         | No               | Three silent-correctness fixes: initialization behind wrappers, trunc_normal_ bounds, pipeline stage coverage                |
 |    |   |           |        |                                                                |
 
 
@@ -217,4 +218,48 @@ This PR improves training monitoring and logging across runs besides some other
 * Add tutorials on Einsum Transformer (Example model integration) and profiling
 
 **Breaking Changes**
-* experiments_root_path is now exposed on an API level
+* experiments_root_path is now exposed on an API level
+
+## PR Initialization and pipeline stage fixes
+
+Three pre-existing bugs, each of which produced a silently wrong model rather than an error.
+
+**1. Weight initialization was skipped entirely for wrapped models.**
+`NamedParameterwiseNormalInitialization` and `Llama3Initializer` stripped only torch.compile's
+`_orig_mod.` prefix from parameter names before matching their regexes. Activation checkpointing and
+FSDP1 insert their own segments (`_checkpoint_wrapped_module.`, `_fsdp_wrapped_module.`), so any
+config that wrapped the model before initializing it matched *no* per-layer regex and kept the
+default initialization. All wrapper prefixes are now normalized away by a shared
+`normalize_parameter_name`.
+
+**2. `Llama3Initializer` injected a single out-of-distribution weight into ~0.3% of tensors.**
+Its `trunc_normal_` calls passed `a=-2, b=2`. torch treats those as *absolute* bounds, but the
+intended standard deviations are 0.02 and smaller, so the bounds sat at +-100 to +-283 sigma. That is
+not merely a no-op: the erf limits of the inverse-transform sampler saturate, its singular edge
+becomes reachable, and the final clamp pins the affected element to exactly the bound - putting one
+weight of magnitude 2.0 into a tensor whose intended scale is 0.007. Measured rate: 12 of 4000
+tensors (0.30%). Bounds are now expressed in standard deviations (`_TRUNCATION_IN_STDS = 3.0`),
+matching the convention the same file already used for the output projection and the one used by the
+llama3 reference implementation. After the fix: 0 of 4000.
+
+This is also the root cause of the intermittent `TestLlama3LikeInitialization` failures. The single
+pinned element dominates the sample variance, which is exactly what the test's `std` assertions
+detect. Measured before the fix: 6 failures in 40 runs (15%); after: 0 in 60. The test's `max`/`min`
+assertions were loosened to the old absolute bound and are now tightened to `3 * std`, so they
+actually detect a stray element instead of tolerating one.
+
+**3. Pipeline stage generation silently dropped modules.**
+`StagesGenerator.get_stages` packed split points greedily against a fixed per-stage weight cap while
+looping exactly `num_virtual_stages` times. Whatever remained when the packing did not fit was never
+assigned to any stage - in practice the output split point, producing a pipeline with no `lm_head` on
+any stage. Uniform per-layer weights always happen to fit, which is why GPT2 never triggered it
+(verified: 0 of ~1000 GPT2 configurations drop anything); any generator with non-uniform per-layer
+cost hits it immediately. Stage packing now uses a partitioner that provably assigns every module to
+exactly one stage, and asking for more stages than split points raises instead of returning empty
+stages.
+
+**Breaking changes:**
+* None in API terms, but item 2 changes the *numerics* of `Llama3Initializer`: truncation now happens
+  at 3 standard deviations instead of an effectively unbounded absolute value. Runs seeded before and
+  after this change will not produce identical weights. This is a fix to the intended behaviour, not
+  a re-tuning.

src/modalities/models/gpt2/llama3_like_initialization.py

@@ -8,8 +8,17 @@
 
 from modalities.models.gpt2.gpt2_model import GPT2LLM
 from modalities.nn.model_initialization.initialization_if import ModelInitializationIF
+from modalities.nn.model_initialization.initialization_routines import normalize_parameter_name
 from modalities.utils.logger_utils import get_logger
 
+# Truncation bounds for trunc_normal_, expressed in standard deviations. torch's trunc_normal_
+# takes `a`/`b` as *absolute* values, so they must be scaled by std. Using an absolute bound that
+# is far outside the distribution (e.g. +-2 with std=0.007, i.e. +-283 sigma) is not merely a no-op:
+# the erf limits of the inverse-transform sampler saturate, its singular edge is reachable, and the
+# final clamp pins the affected element to exactly the bound. That injected one weight of magnitude
+# 2.0 into roughly 0.3% of tensors - a 283 sigma outlier that dominated the tensor's variance.
+_TRUNCATION_IN_STDS = 3.0
+
 logger = get_logger(name="llama3 initialization")
 
 
@@ -47,8 +56,6 @@ def _build_regex_to_init(self, use_weight_tying: bool) -> dict[str, tuple[Callab
                 {
                     "mean": 0.0,
                     "std": 0.02,
-                    "a": -2,
-                    "b": 2,
                 },
             ),
             # final attention projection in attention block
@@ -61,8 +68,6 @@ def _build_regex_to_init(self, use_weight_tying: bool) -> dict[str, tuple[Callab
                         if self.depth_init
                         else 0.02 / math.sqrt(2 * self.num_layers)
                     ),
-                    "a": -2,
-                    "b": 2,
                 },
             ),
             # SwiGLU
@@ -71,8 +76,6 @@ def _build_regex_to_init(self, use_weight_tying: bool) -> dict[str, tuple[Callab
                 {
                     "mean": 0.0,
                     "std": 0.02,
-                    "a": -2,
-                    "b": 2,
                 },
             ),
             r"transformer\.h\.\d+\.mlp\.(V|W_2)\.weight": (
@@ -84,8 +87,6 @@ def _build_regex_to_init(self, use_weight_tying: bool) -> dict[str, tuple[Callab
                         if self.depth_init
                         else 0.02 / math.sqrt(2 * self.num_layers)
                     ),
-                    "a": -2,
-                    "b": 2,
                 },
             ),
         }
@@ -136,10 +137,11 @@ def _init_by_fqn_regex(model: nn.Module, regex_to_init: dict[str, tuple[Callable
                     f"Bias initialization is not allowed for Llama3Initializer. Found bias parameter: {parameter_name}"
                 )
             match_count = 0
+            # Strip FQN modifications introduced by torch.compile, activation checkpointing and
+            # FSDP1 so that the regexes can be written against the plain model. Done once, before
+            # the loop, rather than repeatedly inside it.
+            parameter_name = normalize_parameter_name(parameter_name)
             for weight_regex in regex_to_init.keys():
-                parameter_name = parameter_name.replace(
-                    "_orig_mod.", ""
-                )  # remove FQN modification from torch.compile if present
                 if re.fullmatch(weight_regex, parameter_name):
                     init_fn, arg_dict = regex_to_init[weight_regex]
                     if arg_dict["std"] is not None and callable(arg_dict["std"]):
@@ -154,6 +156,11 @@ def _init_by_fqn_regex(model: nn.Module, regex_to_init: dict[str, tuple[Callable
                                 f"Could not extract layer_id from parameter name {parameter_name} "
                                 "for dynamic std calculation"
                             )
+                    if init_fn is trunc_normal_ and "a" not in arg_dict:
+                        # Bounds are expressed relative to std; see _TRUNCATION_IN_STDS.
+                        arg_dict = arg_dict.copy()
+                        arg_dict["a"] = -_TRUNCATION_IN_STDS * arg_dict["std"]
+                        arg_dict["b"] = _TRUNCATION_IN_STDS * arg_dict["std"]
                     init_fn(p, **arg_dict)
                     match_count += 1
                     hits[weight_regex] += 1

src/modalities/models/parallelism/stages_generator.py

@@ -41,29 +41,22 @@ def get_stages(self, num_layers_per_stage: int, pp_dims: int) -> list[list[str]]
         # The computational weight of the input and output modules are estimated
         # based on the number of layers they correspond to.
         potential_split_points = self._get_potential_split_points()
-        # Calculate the weight per stage based on the total weight and number of stages
-        weight_per_stage = math.ceil(sum(weight for _, weight in potential_split_points) / num_virtual_stages)
-        # pack the stages with the layers
-        next_split_point = 0
-        module_names_per_stage: list[list[str]] = []
-        for _ in range(num_virtual_stages):
-            stage_fqns = []
-            stage_weight = 0
-            while next_split_point < len(potential_split_points):
-                fqns, weight = potential_split_points[next_split_point]
-                if weight > weight_per_stage:
-                    raise ValueError(
-                        f"Weight of {weight} for {fqns} exceeds weight per stage {weight_per_stage}. "
-                        "Please adjust the number of stages or the weight distribution."
-                    )
-                if stage_weight + weight > weight_per_stage:
-                    break
-                stage_fqns.extend(fqns)
-                stage_weight += weight
-                next_split_point += 1
-            module_names_per_stage.append(stage_fqns)
-
-        return module_names_per_stage
+        if num_virtual_stages > len(potential_split_points):
+            raise ValueError(
+                f"Cannot build {num_virtual_stages} pipeline stages from only "
+                f"{len(potential_split_points)} split points. Increase num_layers_per_stage or "
+                f"reduce the pipeline degree."
+            )
+        # Pack the split points into contiguous stages, balancing computational weight.
+        #
+        # This used to pack greedily against a fixed per-stage weight cap, looping exactly
+        # num_virtual_stages times. When the packing did not happen to fit, whatever was left over
+        # was never assigned to any stage and was silently discarded - typically the output split
+        # point, producing a pipeline with no lm_head on any stage. Uniform per-layer weights (as
+        # in GPT2) always happen to fit, which is why this went unnoticed; any generator with
+        # non-uniform weights hits it.
+        groups = _partition_contiguous(potential_split_points, num_parts=num_virtual_stages)
+        return [[fqn for fqns, _ in group for fqn in fqns] for group in groups]
 
     @abstractmethod
     def _get_potential_split_points(self) -> list[tuple[list[str], int]]:
@@ -114,3 +107,103 @@ def _get_potential_split_points(
         ]
 
         return potential_split_points
+
+
+def _greedy_pack(split_points: list[tuple[list[str], int]], weight_cap: int) -> list[list[tuple[list[str], int]]]:
+    """
+    Packs split points left to right into contiguous groups, each at most ``weight_cap`` heavy.
+
+    Unlike a fixed-stage-count loop, this consumes every split point: a group is closed and a new
+    one started whenever the cap would be exceeded.
+
+    Args:
+        split_points (list[tuple[list[str], int]]): The split points with their weights, in order.
+        weight_cap (int): Maximum weight per group. Must be at least the heaviest split point.
+
+    Returns:
+        list[list[tuple[list[str], int]]]: The resulting groups, covering every split point.
+    """
+    groups: list[list[tuple[list[str], int]]] = []
+    current: list[tuple[list[str], int]] = []
+    current_weight = 0
+    for split_point in split_points:
+        weight = split_point[1]
+        if current and current_weight + weight > weight_cap:
+            groups.append(current)
+            current, current_weight = [], 0
+        current.append(split_point)
+        current_weight += weight
+    if current:
+        groups.append(current)
+    return groups
+
+
+def _best_binary_split(group: list[tuple[list[str], int]]) -> int:
+    """
+    Finds the index at which splitting a group minimizes the weight of its heavier half.
+
+    Args:
+        group (list[tuple[list[str], int]]): The group to split, with at least two entries.
+
+    Returns:
+        int: The split index, in ``[1, len(group) - 1]``.
+    """
+    weights = [weight for _, weight in group]
+    total = sum(weights)
+    best_index, best_cost = 1, None
+    prefix = 0
+    for index in range(1, len(group)):
+        prefix += weights[index - 1]
+        cost = max(prefix, total - prefix)
+        if best_cost is None or cost < best_cost:
+            best_index, best_cost = index, cost
+    return best_index
+
+
+def _partition_contiguous(
+    split_points: list[tuple[list[str], int]], num_parts: int
+) -> list[list[tuple[list[str], int]]]:
+    """
+    Partitions split points into exactly ``num_parts`` contiguous, non-empty, balanced groups.
+
+    Finds the smallest per-group weight cap for which a left-to-right greedy pass fits within
+    ``num_parts`` groups (binary search over the cap), then splits the heaviest splittable group
+    until the requested count is reached. Every split point is assigned exactly once.
+
+    Args:
+        split_points (list[tuple[list[str], int]]): The split points with their weights, in order.
+        num_parts (int): The exact number of groups to produce.
+
+    Raises:
+        ValueError: If there are fewer split points than requested groups.
+
+    Returns:
+        list[list[tuple[list[str], int]]]: The groups, covering every split point exactly once.
+    """
+    if num_parts > len(split_points):
+        raise ValueError(f"Cannot partition {len(split_points)} split points into {num_parts} groups.")
+    if num_parts == 1:
+        return [list(split_points)]
+
+    weights = [weight for _, weight in split_points]
+    low, high = max(weights), sum(weights)
+    feasible_cap = high
+    while low <= high:
+        candidate = (low + high) // 2
+        if len(_greedy_pack(split_points, weight_cap=candidate)) <= num_parts:
+            feasible_cap = candidate
+            high = candidate - 1
+        else:
+            low = candidate + 1
+
+    groups = _greedy_pack(split_points, weight_cap=feasible_cap)
+    # The binary search only guarantees "at most num_parts" groups. Split the heaviest splittable
+    # group until the requested count is reached; pipeline parallelism needs exactly this many.
+    while len(groups) < num_parts:
+        splittable = [index for index, group in enumerate(groups) if len(group) > 1]
+        heaviest = max(splittable, key=lambda index: sum(weight for _, weight in groups[index]))
+        group = groups.pop(heaviest)
+        split_index = _best_binary_split(group)
+        groups.insert(heaviest, group[split_index:])
+        groups.insert(heaviest, group[:split_index])
+    return groups

src/modalities/nn/model_initialization/initialization_routines.py

@@ -18,6 +18,34 @@ class MultiDeviceGeneratorPolicy(str, Enum):
     ERROR = "error"
 
 
+# Wrappers that insert themselves into a parameter's fully qualified name without changing which
+# logical parameter it is. The initialization filters are written against the plain model FQNs, so
+# these prefixes are stripped before matching. Without this, applying activation checkpointing (or
+# FSDP1) before initialization would silently prevent every per-layer regex from matching, leaving
+# the model with its default rather than its configured initialization.
+_FQN_WRAPPER_PREFIXES = (
+    "_orig_mod.",  # torch.compile
+    "_checkpoint_wrapped_module.",  # activation checkpointing
+    "_fsdp_wrapped_module.",  # FSDP1
+)
+
+
+def normalize_parameter_name(parameter_name: str) -> str:
+    """
+    Removes wrapper prefixes from a parameter's fully qualified name.
+
+    Args:
+        parameter_name (str): The fully qualified parameter name, possibly containing wrapper
+            segments such as ``_checkpoint_wrapped_module.``.
+
+    Returns:
+        str: The name as it would appear on the unwrapped model.
+    """
+    for prefix in _FQN_WRAPPER_PREFIXES:
+        parameter_name = parameter_name.replace(prefix, "")
+    return parameter_name
+
+
 class PlainInitializationConfig(BaseModel):
     mean: float
     std: Annotated[float, Field(strict=True, ge=0.0)] | str  # can be float or "auto"
@@ -85,9 +113,9 @@ def initialize_in_place(self, model: nn.Module):
         weight_regexes = self.parameter_name_regexes.weights
         bias_regexes = self.parameter_name_regexes.biases or []
         for parameter_name, p in model.named_parameters():
-            parameter_name = parameter_name.replace(
-                "_orig_mod.", ""
-            )  # remove FQN modification from torch.compile if present
+            # Strip FQN modifications introduced by torch.compile, activation checkpointing and
+            # FSDP1 so that the filters can be written against the plain model.
+            parameter_name = normalize_parameter_name(parameter_name)
             for weight_regex in weight_regexes:
                 if re.fullmatch(weight_regex, parameter_name):
                     nn.init.normal_(p, mean=self.mean, std=self.std, generator=self._get_generator(p))