transformers
7bf3c5ec - Make `DtensorShardOperation` work on torch 2.6 and 2.7 (#49657)

Commit
6 days ago
Make `DtensorShardOperation` work on torch 2.6 and 2.7 (#49657) * Make DtensorShardOperation work on torch 2.6 and 2.7 Shard._local_shard_size_and_offset only exists from torch 2.8; 2.6 and 2.7 expose the same computation as _local_shard_size_on_dim(..., return_offset=True). Alias that one when the newer helper is missing, so tensor-parallel loading does not raise AttributeError on the declared torch floor. tests/utils/test_distributed_sharding_utils.py on torch 2.6 goes from 23 failures to the same 3 strided-shard cases torch 2.8 has. Related to #49651 * Resolve the Shard helper in one function instead of patching Shard Per review: pick the right torch API inside _local_shard_size_and_offset and call it from both sites, rather than attaching an alias to Shard.
Author
Parents
Loading