[BUG] Skip None buffers in AutoTP checkpoint loading

[adarshsm]

Bug

Loading.load_buffer() in deepspeed/module_inject/auto_tp.py iterates every key in module._buffers and unconditionally does module._buffers[name].data.is_meta. PyTorch lets a buffer be registered as None — e.g. nn.InstanceNorm1d with the default track_running_stats=False registers running_mean, running_var, and num_batches_tracked as None. Loading a checkpoint into a model containing such a module crashes with:

AttributeError: 'NoneType' object has no attribute 'data'

Repro from #8514 (CPU-only, no GPU or transformers needed):

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

Fix

Skip a buffer when it is None before touching .data — there is nothing to cast to a concrete tensor or copy a checkpoint value into for an unset buffer.

Testing

Added tests/unit/module_inject/test_load_buffer.py:

• test_load_buffer_skips_unset_buffers — an nn.InstanceNorm1d (all-None buffers) goes through Loading.load_buffer without raising. This fails with the reported AttributeError on the code before this change.
• test_load_buffer_still_loads_present_buffers — a normal nn.BatchNorm1d buffer is still copied from state_dict as before.

Ran pytest tests/unit/module_inject/test_load_buffer.py — both pass after the fix, and test_load_buffer_skips_unset_buffers reproduces the reported crash when run against the pre-fix code. Also ran flake8 and yapf --style .style.yapf --diff over both changed files with no findings.

Fixes #8514