FSDP2 raises a KeyError on a tied weight while sharding the model

A language model whose output projection shares its tensor with the input embedding fails to shard. When fully_shard rebuilds parameters as DTensors, the shared tensor appears once rather than twice, the mapping from old parameters to new ones loses the tied name, and preparation dies on a KeyError naming a weight that plainly exists in the model.

A preparation-time failure in FSDP2 where two parameter names alias one tied tensor, so the mapping rebuilt during sharding records the tensor once and a lookup by the second name raises a KeyError.

Sharding replaces every parameter with a new object and needs to know which new object corresponds to which old one. That correspondence is built from tensor identity, and identity is precisely what tying collapses. Two names go in, one entry comes out, and the name that lost the race is the one the error reports.

• The named weight is visibly present in the model, so the error reads as impossible
• The failure happens during preparation, before any training step
• The same model trains without complaint under the previous FSDP implementation
• A warning about weight tying appears above the traceback and is easy to read past

• FSDP2 and fully_shard, which represent sharded parameters as DTensors
• Causal language models that tie the output projection to the input embedding, which is most of them
• Setups where the tied tensors would land in different shard groups, such as the embedding at the root and the head below it
• Optimizer state handling, where the same missing key surfaces on save rather than on prepare

Apply the following checklist to a small reproduction: each box below is a positive signal that you are looking at this exact failure rather than a sibling in the same taxonomy.

• ✓Compare the identity of the output projection tensor with the input embedding tensor. If they are the same object, tying is in play and explains the missing name.
• ✓Read the lines immediately above the traceback for a warning about tied weights; it names the configuration change that resolves it.
• ✓Re-run with tying disabled in the configuration. Success confirms the alias rather than a genuinely absent parameter.

KeyError: 'lm_head.weight'
  File ".../accelerate/accelerator.py", in _prepare_fsdp2
    mapping = {p: new_named_params[n] for n, p in old_named_params.items()}
WARNING: model has tied weights; consider setting tie_word_embeddings=False in the config

Timestamps and exact values vary across runs, but the pattern. An info-level start, an early WARN, an ERROR carrying the symptom. Is the actual fingerprint you should alert on. The Denpex platform flags this combination automatically.

Giving the output projection its own tensor makes the two names refer to two objects, so the mapping has an entry for each and the lookup succeeds. Untie-shard-retie achieves the same thing for a window long enough to build the mapping. Neither changes the mathematics of the model; they change whether the sharding layer can express it.

// Typical pattern:
import torch.optim as optim
optimizer = optim.AdamW(model.parameters(), lr=3e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_training_steps)
for step in range(num_training_steps):
    optimizer.zero_grad()
    loss = model(batch)
    loss.backward()
    optimizer.step()
    scheduler.step()

Adapt the snippet to your framework. The same pattern holds for PyTorch Lightning, Hugging Face Trainer, DeepSpeed, Megatron-LM, and vLLM training wrappers. Where the wrapper exposes a config flag (for examplelr_scheduler_type in Trainer), prefer the flag over the imperative API to keep the schedule declarative and reproducible.