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

[le1nux]

Three pre-existing bugs found while working on #459. Each produced a silently wrong model rather than an error, so none of them would show up as a failure — only as a model that trains slightly worse than configured.

Branched from main and kept deliberately small; independent of #459.

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:

parameter_name = parameter_name.replace("_orig_mod.", "")  # torch.compile only

Activation checkpointing and FSDP1 insert their own segments (_checkpoint_wrapped_module., _fsdp_wrapped_module.). So for any config that wraps the model before initializing it — which is the normal ordering, since activation checkpointing must precede FSDP2 wrapping and initialization must follow it — every per-layer regex silently failed to match and the model kept its default nn.Linear initialization.

A shared normalize_parameter_name now strips all wrapper prefixes.

2. Llama3Initializer injected an out-of-distribution weight into ~0.3% of tensors

Its trunc_normal_ calls passed a=-2, b=2. torch treats these as absolute bounds, but the intended standard deviations are 0.02 and smaller — so the bounds sat at ±100σ to ±283σ.

That is not merely a no-op. With bounds that far out, the erf limits of the inverse-transform sampler saturate to ±1, its singular edge becomes reachable, and the final clamp_ pins the affected element to exactly the bound. The result is one weight of magnitude 2.0 in a tensor whose intended scale is 0.007 — a 283σ outlier.

The signature is unmistakable: affected tensors have |w|max == 2.00000 exactly, while the remaining 65535 elements are a perfect normal (0.28% beyond 3σ, versus 0.27% expected). And the sample variance matches the prediction for "correct normal plus one element at ±2.0" to four decimals:

target std 0.00707 + one element at |2.0| -> predicted 0.01054   observed 0.01053
target std 0.02000 + one element at |2.0| -> predicted 0.02147   observed 0.02151

Measured rate:

std=0.00707  absolute +-2 (before)  12/4000 = 0.30%   worst sample std 0.01057
std=0.02000  absolute +-2 (before)  12/4000 = 0.30%   worst sample std 0.02157
std=0.00707  relative +-3*std (after)  0/4000 = 0.00%
std=0.02000  relative +-3*std (after)  0/4000 = 0.00%

Bounds are now expressed in standard deviations (_TRUNCATION_IN_STDS = 3.0). That matches both the convention this same file already used for the output projection (3 / sqrt(n_embd), i.e. 3 * std) and the llama3 reference in TorchTitan (a=-3*s, b=3*s) — which is what makes me confident the ±2 was a slip rather than intent.

This is also the root cause of the flaky TestLlama3LikeInitialization

The pinned element dominates the tensor's sample variance, which is precisely what the test's std assertions detect. So the test was correct and the code was wrong:

	failures
before	6 in 40 runs (15%)
after	0 in 60 runs

The 15% rate is itself predicted by the 0.30% per-tensor rate: with ~56 checked tensors across a run, 1 - 0.997^56 = 15.5%.

The test's max/min assertions had been set to the old absolute ±2 bound, which is why they tolerated a 2.0 weight instead of catching it. They now assert 3 * std, so a stray element fails the test.

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. Anything left over when the packing didn't fit was never assigned to any stage:

MISSING: {'transformer.lm_head_norm', 'transformer.lm_head'}

i.e. a pipeline with no output layer on any stage.

Uniform per-layer weights always happen to fit, which is why GPT-2 never triggered this — I verified 0 of ~1000 GPT-2 configurations drop anything. Any generator with non-uniform per-layer cost (e.g. a hybrid stack where an MoE layer costs 3× a Mamba layer) hits it immediately.

Packing now uses a partitioner that provably assigns every module to exactly one stage: binary-search the smallest feasible per-stage weight cap, then split the heaviest group until the requested stage count is reached. Requesting more stages than split points now raises instead of returning empty stages.

Testing

New: tests/models/parallelism/test_stages_generator.py, tests/nn/model_initialization/test_fqn_normalization.py.

Both suites were checked against the unfixed code to confirm they actually catch the bugs — 5 stage tests fail on the old packer, including the one asserting lm_head survives, while the GPT-2 cases pass either way (confirming GPT-2 was unaffected).

Suite	Result
tests/models, nn, config, checkpoint strategies, optimizer	129 passed, 1 skipped
test_initialization_fsdpx.py + test_initialization_fsdp1.py (2 GPUs)	26 passed, 20 skipped
TestLlama3LikeInitialization × 60 consecutive runs	60 passed
isort / black / ruff	clean

One behavioural note for reviewers

Item 2 changes the numerics of Llama3Initializer: truncation now happens at 3 standard deviations instead of at an effectively unbounded absolute value. Seeded runs from before and after this change will not produce identical weights. I consider that a fix to the intended behaviour rather than a re-tuning — the old bounds were unreachable by design and actively harmful when reached — but it is a real change to initialization and worth a second opinion.

Item 3 overlaps slightly with #459, which works around the same packer bug by overriding get_stages in its own generator. Once this lands, that override becomes redundant and can be deleted. #459 also contains the same fix as item 1; whichever merges first, the other needs a trivial conflict resolution.

🤖 Generated with Claude Code