AutoTP loads module weights by hardcoded class name, so 219 of 233 transformers norm classes are silently left on meta

[alanhuangyoo]

Summary

Loading.is_load_module decides whether AutoTP loads a module's weights from the checkpoint, and it decides by class name against a hardcoded list of 26. A norm class not on that list never has its weights loaded: on the meta-device path its parameters stay on meta.

219 of the 233 norm classes in transformers 5.16.1 are not on the list, including every Gemma generation, every GLM, Granite, Olmo, MiniMax, Nemotron and Zamba.

Demonstration

Two norm classes with identical structure — same base, same weight parameter, same forward — differing only in name. Model built under torch.device("meta"), then run through the AutoTP replacement with a state dict that gives both norms 7.0:

RESULT checkpoint says 7.0 for both norms
  LlamaRMSNorm  (on the list):  value=7.0
  GemmaRMSNorm  (not on it):    STILL ON META

The gate itself, on classes taken from transformers:

  LlamaRMSNorm             is_load_module=True
  Qwen3_5RMSNormGated      is_load_module=True
  GemmaRMSNorm             is_load_module=False
  Gemma3RMSNorm            is_load_module=False
  CohereLayerNorm          is_load_module=False
  DeepseekV32RMSNorm       is_load_module=False
  MiniMaxRMSNorm           is_load_module=False
  GraniteRMSNorm           is_load_module=False

Why nothing catches it

# deepspeed/module_inject/auto_tp.py
if Loading.is_load_module(child) and self.state_dict is not None:
    if any(checking_key in item for item in self.state_dict):
        Loading.load(child, self.state_dict, checking_key, self.mp_group)
    else:
        continue

Loading.load is the only thing that materializes a meta parameter here, and it is behind the gate. A norm has no children, so falling through to the recursion reaches nothing. The one leftover-meta case that is handled downstream is tied embeddings:

# deepspeed/module_inject/replace_module.py
if embedding_weight is not None and hasattr(module, "lm_head") and ... module.lm_head.weight.is_meta:
    module.lm_head.weight = embedding_weight

which is a targeted fix for one module rather than a general sweep — if there were a general one, that line would not be needed.

What is missing

Restricting to families people actually run, 47 classes:

Gemma2RMSNorm, Gemma3RMSNorm, Gemma3nRMSNorm, Gemma4RMSNorm, Gemma4UnifiedRMSNorm, GemmaRMSNorm,
Glm4MoeLiteRMSNorm, Glm4MoeRMSNorm, Glm4RMSNorm, Glm4vMoeRMSNorm, Glm4vMoeTextRMSNorm, Glm4vRMSNorm,
Glm5NextRMSNorm, Glm5NextTextRMSNorm, Glm5NextTextRMSNormGated, Glm5NextTextUnweightedRMSNorm,
GlmImageRMSNorm, GlmMoeDsaRMSNorm, GlmOcrRMSNorm, GlmRMSNorm,
Granite4VisionTextRMSNorm, GraniteMoeHybridRMSNorm, GraniteMoeHybridRMSNormGated, GraniteMoeRMSNorm,
GraniteMoeSWARMSNorm, GraniteMoeSharedRMSNorm, GraniteRMSNorm, GraniteSWARMSNorm,
MiniCPM3RMSNorm, MiniMaxM2RMSNorm, MiniMaxM3VLRMSNorm, MiniMaxRMSNorm, NemotronHRMSNorm,
Olmo2RMSNorm, Olmo3RMSNorm, OlmoHybridRMSNorm, OlmoHybridRMSNormGated, OlmoLayerNorm, OlmoeRMSNorm,
Phi4MultimodalRMSNorm, RecurrentGemmaRMSNorm, T5Gemma2RMSNorm, T5GemmaRMSNorm, VaultGemmaRMSNorm,
Zamba2RMSNorm, Zamba2RMSNormGated, ZambaRMSNorm

GlmMoeDsaRMSNorm is the norm of inference-optimization/GLM-5.2-0.8B-A0.8B, one of the models used in this repo's own testing.

Where this leaves the maintenance model

#8306 added four names for Qwen3.5, which is the right fix under the current design and also shows the problem with it: the list has to grow by hand for every model family and every generation of it, and when it does not, the failure is silent rather than an error. transformers gains norm classes faster than the list does.

Two directions, and I would rather have your call than pick one:

1. Decide structurally instead of by name. The question the gate is really asking is "is this a leaf module whose parameters the checkpoint has and nothing else will replace". That is answerable from the module and the state dict — no list. It is the bigger change and needs care about what should not be loaded here.
2. Fail loudly instead of silently. Keep the list, but after replacement, check for parameters still on meta and raise naming them. That does not fix any model, but it converts a silent wrong-weights run into an error that says which class to add — which is what the current design needs to be usable.

Doing 2 is small and would have surfaced this immediately; 1 removes the maintenance entirely. I am happy to write either.

What I did not verify

I demonstrated the gate and the meta parameter surviving replacement, not an end-to-end wrong generation from a real Gemma checkpoint — I do not have that setup here. If someone has an AutoTP inference run against a checkpoint for any family above, the quick check is whether any parameter is still .is_meta after init_inference returns.

Reproduction

import torch, deepspeed
from torch import nn
from deepspeed.module_inject.auto_tp import AutoTP
from deepspeed.module_inject.layers import set_autotp_mode

def norm_cls(name):
    def __init__(self, h=8):
        nn.Module.__init__(self); self.weight = nn.Parameter(torch.ones(h))
    return type(name, (nn.Module,), {"__init__": __init__})

class M(nn.Module):
    def __init__(self, h=8):
        super().__init__()
        self.lin = nn.Linear(h, h, bias=False)
        self.listed_norm = norm_cls("LlamaRMSNorm")(h)
        self.unlisted_norm = norm_cls("GemmaRMSNorm")(h)

with torch.device("meta"):
    model = M()
sd = {"model.lin.weight": torch.eye(8),
      "model.listed_norm.weight": torch.full((8,), 7.0),
      "model.unlisted_norm.weight": torch.full((8,), 7.0)}

set_autotp_mode(training=True)
autotp = AutoTP(module=model, all_reduce_linears=(), prefix="model", state_dict=sd,
                linear_layer_setting=(nn.Linear, nn.Embedding), orig_layer_impl=None, training_mode=True)
autotp.mp_size, autotp.mp_group = 1, None
autotp.update_linear_policies()
autotp._replace_module(model)

print(model.listed_norm.weight.is_meta, model.unlisted_norm.weight.is_meta)  # False True