Skip to content

fix: skip None buffers in AutoTP Loading.load_buffer (checkpoint load with nn.InstanceNorm1d) - #8603

Open
SIDDARTHAREDDY8 wants to merge 1 commit into
deepspeedai:masterfrom
SIDDARTHAREDDY8:bhai-blitz-8-deepspeed-1
Open

SIDDARTHAREDDY8 wants to merge 1 commit into
deepspeedai:masterfrom
SIDDARTHAREDDY8:bhai-blitz-8-deepspeed-1

Conversation

@SIDDARTHAREDDY8

Copy link
Copy Markdown

Fixes #8514.

Summary
Loading a checkpoint via replace_module/AutoTP into a model containing nn.InstanceNorm1d raised
AttributeError: 'NoneType' object has no attribute 'data' in Loading.load_buffer
(deepspeed/module_inject/auto_tp.py:153). InstanceNorm's running-stat buffers are None until first
use, so module._buffers[name].data.is_meta blows up even though there is nothing to load for them.
The fix skips None buffers in the loop, matching the checkpoint=None walk behavior.

Test notes

  • Reproduced the exact AttributeError from the issue on master's load_buffer with an nn.InstanceNorm1d
    module and a saved state dict; post-fix the walk completes and the None buffers are preserved as None.
  • Regression-checked that modules with real buffers (e.g. nn.BatchNorm1d) still load checkpoint values
    correctly and that the meta-tensor materialization path is unchanged.
  • Full deepspeed package import requires compiled ops, so the check ran the exact shipped function body
    (AST-extracted from the patched file) under torch 2.14.0+cpu. No GPU available here.

… with nn.InstanceNorm1d)

Signed-off-by: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant