[BUG] AutoTP checkpoint load raises AttributeError on a model containing nn.InstanceNorm1d #8514

Describe the bug

Loading a checkpoint into a model that has an nn.InstanceNorm1d child raises AttributeError: 'NoneType' object has no attribute 'data' in Loading.load_buffer (deepspeed/module_inject/auto_tp.py:153). The same model with checkpoint=None walks through fine.

To Reproduce

import torch, torch.nn as nn
from deepspeed.module_inject.replace_module import replace_module

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = nn.InstanceNorm1d(4)
        self.fc = nn.Linear(4, 4)

torch.save(Net().state_dict(), "/tmp/ckpt.pt")
noop = lambda child, policy, layer_id, prefix='', state_dict=None: child

replace_module(Net(), nn.Linear, noop, None, checkpoint=None)             # returns
replace_module(Net(), nn.Linear, noop, None, checkpoint="/tmp/ckpt.pt")   # raises

python repro.py, on master 8bdd8f19326bfc22d87bfcf509ebff114384d13d. torch is the only dependency, transformers is not needed. noop is there so the walk runs without a real injection policy.

  File "deepspeed/module_inject/replace_module.py", line 705, in _replace_module
    Loading.load_buffer(child, state_dict, checking_key)
  File "deepspeed/module_inject/auto_tp.py", line 153, in load_buffer
    if module._buffers[name].data.is_meta:
AttributeError: 'NoneType' object has no attribute 'data'

For context on what the two sides look like: Net().norm._buffers is {'running_mean': None, 'running_var': None, 'num_batches_tracked': None}, and the saved checkpoint holds only fc.weight and fc.bias.

I drove replace_module directly because I have no GPU box here for a full init_inference run. AutoTP._replace_module at auto_tp.py:749 makes the same Loading.load_buffer call, so I would expect it on that path too, but I have not run it.

Expected behavior

The walk finishes, the way it does with checkpoint=None. There is nothing under norm. in the checkpoint to load.

ds_report output

torch install path ............... ['/usr/local/lib/python3.12/site-packages/torch']
torch version .................... 2.14.0+cpu
deepspeed install path ........... ['/work/deepspeed']
deepspeed info ................... 0.19.6, [none], [none]
deepspeed wheel compiled w. ...... torch 0.0

System info

• OS: Debian 13 (python:3.12-slim)
• GPU count and types: none, CPU only
• Python 3.12.14
• torch 2.14.0+cpu
• transformers: not installed
• DeepSpeed: master at 8bdd8f19326bfc22d87bfcf509ebff114384d13d

Docker context

python:3.12-slim with torch 2.14.0+cpu installed, nothing else.