From 3e1a5bac72c5b5e3d9592196a96022b99dda9d79 Mon Sep 17 00:00:00 2001 From: Max Luebbering <2804731+le1nux@users.noreply.github.com> Date: Thu, 30 Jul 2026 15:39:57 +0200 Subject: [PATCH] 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) --- CHANGELOG_DEV.md | 47 +++++- .../models/gpt2/llama3_like_initialization.py | 29 ++-- .../models/parallelism/stages_generator.py | 139 +++++++++++++++--- .../initialization_routines.py | 34 ++++- tests/models/parallelism/__init__.py | 0 .../parallelism/test_stages_generator.py | 121 +++++++++++++++ .../test_fqn_normalization.py | 70 +++++++++ tests/test_initialization_fsdpx.py | 38 +++-- 8 files changed, 419 insertions(+), 59 deletions(-) create mode 100644 tests/models/parallelism/__init__.py create mode 100644 tests/models/parallelism/test_stages_generator.py create mode 100644 tests/nn/model_initialization/test_fqn_normalization.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 43d0c6e2d..ee797ef83 100644 --- a/CHANGELOG_DEV.md +++ b/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 \ No newline at end of file +* 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. diff --git a/src/modalities/models/gpt2/llama3_like_initialization.py b/src/modalities/models/gpt2/llama3_like_initialization.py index ea6870b16..802dd5e6d 100644 --- a/src/modalities/models/gpt2/llama3_like_initialization.py +++ b/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 diff --git a/src/modalities/models/parallelism/stages_generator.py b/src/modalities/models/parallelism/stages_generator.py index eeb78f01f..0899bd5a8 100644 --- a/src/modalities/models/parallelism/stages_generator.py +++ b/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 diff --git a/src/modalities/nn/model_initialization/initialization_routines.py b/src/modalities/nn/model_initialization/initialization_routines.py index 1f785f562..d41eb211a 100644 --- a/src/modalities/nn/model_initialization/initialization_routines.py +++ b/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)) diff --git a/tests/models/parallelism/__init__.py b/tests/models/parallelism/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/models/parallelism/test_stages_generator.py b/tests/models/parallelism/test_stages_generator.py new file mode 100644 index 000000000..532abe0cf --- /dev/null +++ b/tests/models/parallelism/test_stages_generator.py @@ -0,0 +1,121 @@ +"""Tests that pipeline stage generation assigns every module to exactly one stage. + +The generator used to pack split points greedily against a fixed per-stage weight cap while looping +exactly ``num_virtual_stages`` times. Anything left over when the packing did not fit was silently +discarded - in practice 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, so this only surfaced for a +generator with non-uniform weights. +""" + +import pytest + +from modalities.models.parallelism.stages_generator import GPT2LLMStagesGenerator, StagesGenerator + + +class _NonUniformStagesGenerator(StagesGenerator): + """A generator whose layers have differing computational cost, e.g. a hybrid MoE stack.""" + + def __init__(self, layer_weights: list[int], input_layer_equivalence: int = 2, output_layer_equivalence: int = 2): + super().__init__( + num_model_layers=len(layer_weights), + input_layer_equivalence=input_layer_equivalence, + output_layer_equivalence=output_layer_equivalence, + ) + self._layer_weights = layer_weights + + def _get_potential_split_points(self) -> list[tuple[list[str], int]]: + return [ + (["transformer.wte"], self._input_layer_equivalence), + *[([f"transformer.h.{i}"], w) for i, w in enumerate(self._layer_weights)], + (["transformer.lm_head_norm", "transformer.lm_head"], self._output_layer_equivalence), + ] + + +def _all_modules(generator: StagesGenerator) -> set[str]: + return {fqn for fqns, _ in generator._get_potential_split_points() for fqn in fqns} + + +@pytest.mark.parametrize( + "layer_weights,num_layers_per_stage,pp_dims", + [ + # Expensive layers first: the greedy pack exhausts its budget early and used to drop the tail. + ([3, 3, 3, 3, 2, 2, 2, 2], 6, 2), + ([2, 2, 2, 2, 3, 3, 3, 3], 6, 2), + ([3, 2, 3, 2, 3, 1, 3, 2], 6, 2), + ([3, 2, 3, 2, 3, 1, 3, 2], 3, 4), + # A Nemotron-3 Nano-shaped stack: 23 Mamba (2), 23 MoE (3), 6 attention (1). + ([2, 3] * 23 + [1] * 6, 14, 4), + ([2, 3] * 23 + [1] * 6, 28, 2), + ], +) +def test_every_module_is_assigned_exactly_once(layer_weights, num_layers_per_stage, pp_dims): + generator = _NonUniformStagesGenerator(layer_weights) + stages = generator.get_stages(num_layers_per_stage=num_layers_per_stage, pp_dims=pp_dims) + flat = [fqn for stage in stages for fqn in stage] + + assert set(flat) == _all_modules(generator), "a module was dropped" + assert len(flat) == len(set(flat)), "a module was assigned to more than one stage" + assert all(stage for stage in stages), "an empty pipeline stage was produced" + assert len(stages) % pp_dims == 0 + + +def test_output_layer_is_never_dropped(): + # The concrete failure: transformer.lm_head vanished from every stage. + generator = _NonUniformStagesGenerator([3, 3, 3, 3, 2, 2, 2, 2]) + stages = generator.get_stages(num_layers_per_stage=6, pp_dims=2) + assert "transformer.lm_head" in stages[-1] + assert "transformer.lm_head_norm" in stages[-1] + assert "transformer.wte" in stages[0] + + +def test_stages_preserve_model_order(): + generator = _NonUniformStagesGenerator([2, 3, 2, 3, 1, 2, 3, 2]) + stages = generator.get_stages(num_layers_per_stage=6, pp_dims=2) + flat = [fqn for stage in stages for fqn in stage] + layer_indices = [int(f.split(".")[-1]) for f in flat if f.startswith("transformer.h.")] + assert layer_indices == sorted(layer_indices) + + +def test_stage_weights_are_near_balanced(): + layer_weights = [2, 3] * 23 + [1] * 6 + generator = _NonUniformStagesGenerator(layer_weights) + weight_by_fqn = {fqns[0]: w for fqns, w in generator._get_potential_split_points()} + stages = generator.get_stages(num_layers_per_stage=14, pp_dims=4) + + stage_weights = [sum(weight_by_fqn[f] for f in stage if f in weight_by_fqn) for stage in stages] + ideal = sum(weight_by_fqn.values()) / len(stages) + # The slowest stage sets pipeline throughput, so the imbalance must stay small. + assert max(stage_weights) <= ideal * 1.15 + + +def test_rejects_more_stages_than_split_points(): + generator = _NonUniformStagesGenerator([1, 1]) + with pytest.raises(ValueError, match="Cannot build"): + generator.get_stages(num_layers_per_stage=1, pp_dims=6) + + +# -------------------------------------------------------------------------------------------- +# GPT2 behaviour must be unchanged: it was never affected, and its stages should stay complete. +# -------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "num_model_layers,num_layers_per_stage,pp_dims", + [(12, 7, 2), (12, 14, 1), (4, 3, 2), (24, 13, 2), (8, 5, 2)], +) +def test_gpt2_stages_remain_complete(num_model_layers, num_layers_per_stage, pp_dims): + generator = GPT2LLMStagesGenerator(num_model_layers=num_model_layers) + stages = generator.get_stages(num_layers_per_stage=num_layers_per_stage, pp_dims=pp_dims) + flat = [fqn for stage in stages for fqn in stage] + + assert set(flat) == _all_modules(generator) + assert len(flat) == len(set(flat)) + assert len(stages) % pp_dims == 0 + layer_indices = [int(f.split(".")[-1]) for f in flat if f.startswith("transformer.h.")] + assert layer_indices == sorted(layer_indices) + + +def test_non_divisible_stage_count_still_raises(): + generator = GPT2LLMStagesGenerator(num_model_layers=12) + with pytest.raises(ValueError, match="not divisible by parallel dimensions"): + generator.get_stages(num_layers_per_stage=5, pp_dims=2) diff --git a/tests/nn/model_initialization/test_fqn_normalization.py b/tests/nn/model_initialization/test_fqn_normalization.py new file mode 100644 index 000000000..c4b51e154 --- /dev/null +++ b/tests/nn/model_initialization/test_fqn_normalization.py @@ -0,0 +1,70 @@ +"""Tests that initialization filters still match when the model is wrapped. + +Activation checkpointing, torch.compile and FSDP1 all insert segments into a parameter's fully +qualified name. The initialization filters are written against the plain model, so those segments +have to be stripped before matching. Without that, applying activation checkpointing before weight +initialization makes every per-layer regex fail and the model silently keeps its default +initialization - it trains, just not as configured. +""" + +import pytest +import torch +import torch.nn as nn + +from modalities.nn.model_initialization.initialization_routines import InitializationRoutines, normalize_parameter_name +from modalities.nn.model_initialization.parameter_name_filters import RegexFilter + + +@pytest.mark.parametrize( + "wrapped,plain", + [ + ("transformer.h.0._checkpoint_wrapped_module.attn.c_proj.weight", "transformer.h.0.attn.c_proj.weight"), + ("_orig_mod.transformer.wte.weight", "transformer.wte.weight"), + ("_fsdp_wrapped_module.transformer.lm_head.weight", "transformer.lm_head.weight"), + # Several wrappers can be layered (compile over activation checkpointing over FSDP1). + ( + "_orig_mod.transformer.h.3._checkpoint_wrapped_module.mlp.W.weight", + "transformer.h.3.mlp.W.weight", + ), + # An unwrapped name must pass through untouched. + ("transformer.h.1.attn.q_attn.weight", "transformer.h.1.attn.q_attn.weight"), + ], +) +def test_normalize_parameter_name_strips_wrapper_segments(wrapped, plain): + assert normalize_parameter_name(wrapped) == plain + + +class _WrappedBlock(nn.Module): + """Mimics how activation checkpointing renames a submodule's parameters.""" + + def __init__(self): + super().__init__() + self._checkpoint_wrapped_module = nn.Linear(64, 64, bias=False) + + +class _Model(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.ModuleDict({"h": nn.ModuleDict({"0": _WrappedBlock()})}) + + @property + def wrapped_weight(self) -> nn.Parameter: + return self.transformer["h"]["0"]._checkpoint_wrapped_module.weight + + +def test_initializer_matches_parameters_behind_a_wrapper(): + model = _Model() + # The filter is written against the plain FQN, which is the whole point. + regex_filter = RegexFilter(weights=[r"transformer\.h\.\d+\.weight"]) + assert any("_checkpoint_wrapped_module" in name for name, _ in model.named_parameters()) + + std = 0.02 + initializer = InitializationRoutines.get_plain_initialization( + mean=0.0, std=std, parameter_name_regexes=regex_filter, seed=42 + ) + with torch.no_grad(): + model.wrapped_weight.fill_(123.0) + initializer.initialize_in_place(model) + + assert model.wrapped_weight.std().item() == pytest.approx(std, rel=0.1) + assert model.wrapped_weight.abs().max().item() < 1.0, "parameter was not re-initialized" diff --git a/tests/test_initialization_fsdpx.py b/tests/test_initialization_fsdpx.py index f1eb4210a..31471a234 100644 --- a/tests/test_initialization_fsdpx.py +++ b/tests/test_initialization_fsdpx.py @@ -548,36 +548,32 @@ def _test_qkv_proj(self, gpt2_block: GPT2Block): layers = (gpt2_block.attn.q_attn, gpt2_block.attn.k_attn, gpt2_block.attn.v_attn) for layer in layers: assert layer.weight.std().detach().cpu() == pytest.approx(0.02, abs=1e-3) - assert layer.weight.max().detach().cpu() <= 2 - assert layer.weight.min().detach().cpu() >= -2 + # trunc_normal_ bounds are 3 * std; a weight outside them means an out-of-distribution + # element slipped in (see _TRUNCATION_IN_STDS in llama3_like_initialization.py). + assert layer.weight.max().detach().cpu() <= 3 * 0.02 + assert layer.weight.min().detach().cpu() >= -3 * 0.02 assert layer.weight.mean().detach().cpu() == pytest.approx(0, abs=1e-3) def _test_c_proj(self, gpt2_block: GPT2Block, depth_init: bool, n_layer: int, layer_id: int): layer = gpt2_block.attn.c_proj - if depth_init: - assert layer.weight.std().detach().cpu() == pytest.approx(0.02 / math.sqrt(2 * (layer_id + 1)), abs=1e-3) - else: - assert layer.weight.std().detach().cpu() == pytest.approx(0.02 / math.sqrt(2 * n_layer), abs=1e-3) - - assert layer.weight.max().detach().cpu() <= 2 - assert layer.weight.min().detach().cpu() >= -2 + expected_std = 0.02 / math.sqrt(2 * (layer_id + 1) if depth_init else 2 * n_layer) + assert layer.weight.std().detach().cpu() == pytest.approx(expected_std, abs=1e-3) + # Bounds are 3 * std; anything outside means an out-of-distribution element. + assert layer.weight.max().detach().cpu() <= 3 * expected_std + assert layer.weight.min().detach().cpu() >= -3 * expected_std assert layer.weight.mean().detach().cpu() == pytest.approx(0, abs=1e-3) def _test_swiglu_proj(self, gpt2_block: GPT2Block, depth_init: bool, n_layer: int, layer_id: int): - layers = (gpt2_block.mlp.V, gpt2_block.mlp.W_2) - for layer in layers: - if depth_init: - assert layer.weight.std().detach().cpu() == pytest.approx( - 0.02 / math.sqrt(2 * (layer_id + 1)), abs=1e-3 - ) - else: - assert layer.weight.std().detach().cpu() == pytest.approx(0.02 / math.sqrt(2 * n_layer), abs=1e-3) - assert layer.weight.max().detach().cpu() <= 2 - assert layer.weight.min().detach().cpu() >= -2 + expected_std = 0.02 / math.sqrt(2 * (layer_id + 1) if depth_init else 2 * n_layer) + for layer in (gpt2_block.mlp.V, gpt2_block.mlp.W_2): + assert layer.weight.std().detach().cpu() == pytest.approx(expected_std, abs=1e-3) + # Bounds are 3 * std; anything outside means an out-of-distribution element. + assert layer.weight.max().detach().cpu() <= 3 * expected_std + assert layer.weight.min().detach().cpu() >= -3 * expected_std assert layer.weight.mean().detach().cpu() == pytest.approx(0, abs=1e-3) layer = gpt2_block.mlp.W assert layer.weight.std().detach().cpu() == pytest.approx(0.02, abs=1e-3) - assert layer.weight.max().detach().cpu() <= 2 - assert layer.weight.min().detach().cpu() >= -2 + assert layer.weight.max().detach().cpu() <= 3 * 0.02 + assert layer.weight.min().detach().cpu() >= -3 * 0.02 assert layer.weight.mean().detach().cpu() == pytest.approx(0, abs=1e-3)