Fix DTensor double-backward patch for torch >= 2.13 - #447
Open
luciaquirke wants to merge 1 commit into
Open
Conversation
apply_dtensor_patch replaces Redistribute.backward with a version that reads
ctx.backward_dtype and ctx.original_dtype. torch's Redistribute.forward stored
those up to 2.12; from 2.13 it routes dtypes through a _DtypeConfig and stores
ctx.bwd_op_dtype / ctx.bwd_out_dtype instead. The old names are therefore absent
and the first backward raises
AttributeError: 'RedistributeBackward' object has no attribute 'backward_dtype'
which makes fsdp=True unusable for any model. Reproduced on torch 2.13.0+cu126
with gpt2 and nproc_per_node=2, so it is not size- or architecture-specific.
Read the dtypes through a helper that accepts either spelling, so the patch works
on both. With the fix the same gpt2 + FSDP run trains normally.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
apply_dtensor_patchreplacesRedistribute.backwardwith a version that readsctx.backward_dtypeandctx.original_dtype.torch's
Redistribute.forwardstored those attributes up to 2.12. From 2.13 it routes dtypes through a_DtypeConfigand storesctx.bwd_op_dtype/ctx.bwd_out_dtypeinstead, so the old names are absent and the first backward raises:This makes
fsdp=Trueunusable for any model.Reproduction
torch 2.13.0+cu126,
gpt2,nproc_per_node=2,fsdp: true— fails on the first backward. Not size- or architecture-specific: it surfaced while trying to train a 7B, but a 124M model fails identically, which is what isolated it to the patch rather than to model scale.Fix
Read the dtypes through a small helper that accepts either spelling, so the patch keeps working on torch <= 2.12 and on >= 2.13.
With the fix applied, the same gpt2 + FSDP run trains normally.
🤖 Generated with Claude Code
https://claude.ai/code/session_0167XKnE2yoCE162n3jHvxC1