Repository: pytorch/pytorch

Issue:
(bug)(torch2.0/datapipes) Potentially backwards incompatible change with DataLoader and `is_shardable` Datapipes
### 🐛 Describe the bug

[PR-94095](https://github.com/pytorch/pytorch/pull/94095) breaks BC for custom `DataPipe` implementations that implement `is_shardable()` and `apply_sharding()`. 

The BC breaking change is that the `apply_sharding()` method is not called on the datapipe from the `DataLoader` unless the datapipe is a subclas of `_ShardingIterDataPipe`. Therefore,  on custom datapipes that implement `is_shardable()` and `apply_sharding()` you end up getting duplicate elements in the batch when using with a dataloader with `num_workers > 1`.

(please correct me if I'm wrong but) I didn't see any mention of this BC breaking change in the release notes for:

1. `torchdata-0.6.0`: https://github.com/pytorch/data/releases
2. `torch-2.0`: https://github.com/pytorch/pytorch/releases/tag/v2.0.0

Since the PR introduces a "private" (`_ShardingIterDataPipe`) it makes it unnatural to subclass from existing custom sharded datapipes. I believe that the intent of the PR was to have users use `ShardingFilterIterDataPipe` (e.g. chain the source pipe with `.sharding_filter()`) but this may be unfeasible due to:

1. Custom sharding strategies where `sharding_filter`'s modulo-n style sharding does not suffice.
2. Need to change all the call-sites to have callers additionally call `my_datapipe().sharding_filter()` as `my_datapipe` is not longer implicitly sharded.

I'm wondering what the best next steps are? Subclass `_ShardingIterDataPipe`? or fix the BC break (we could use a `@runtime_checkable` as:

```
@runtime_checkable
class Shardable(Protocol):
    """
    Indicates whether a datapipe is shardable. A shardable datapipe distributes (shards) the data records
    amongst the dataloader worker instances in such a way that the same record is not processed twice
    when the dataloader ``num_workers > 1``.

    .. note::
      A shardable datapipe must implement BOTH the indicator method ``is_shardable()``
      AND the sharding logic ``apply_sharding(num_instances, instance_id)``.

    Use with ``isinstance()`` to check whether a datapipe is shardable

    .. doctest::

     >>> from torchdata.datapipes.iter import IterableWrapper
     >>> from ape.datapipes import Shardable
     >>> dp = IterableWrapper([1,2,3]).sharding_filter()
     >>> isinstance(dp, Shardable)
     True

    """

    def is_shardable(self) -> bool:  # pragma: no cover
        ...

    def apply_sharding(
        self, num_instances: int, instance_id: int
    ) -> None:  # pragma: no cover
        ...

```



### Versions

 ubuntu@ip-10-2-87-95% python collect_env.py                                                                   ~/tmp
Collecting environment information...
PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.26.0
Libc version: glibc-2.31

Python version: 3.9.15 (main, Feb 13 2023, 21:22:03)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-1015-aws-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 510.73.08
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          96
On-line CPU(s) list:             0-95
Thread(s) per core:              2
Core(s) per socket:              24
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel(R) Xeon(R) Platinum 8275CL CPU @ 3.00GHz
Stepping:                        7
CPU MHz:                         2999.998
BogoMIPS:                        5999.99
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       1.5 MiB
L1i cache:                       1.5 MiB
L2 cache:                        48 MiB
L3 cache:                        71.5 MiB
NUMA node0 CPU(s):               0-23,48-71
NUMA node1 CPU(s):               24-47,72-95
Vulnerability Itlb multihit:     KVM: Mitigation: VMX unsupported
Vulnerability L1tf:              Mitigation; PTE Inversion
Vulnerability Mds:               Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Meltdown:          Mitigation; PTI
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, STIBP disabled, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves ida arat pku ospke

Versions of relevant libraries:
[pip3] mypy==0.991
[pip3] mypy-boto3-batch==1.26.34
[pip3] mypy-boto3-ec2==1.26.71
[pip3] mypy-boto3-iam==1.26.62
[pip3] mypy-boto3-s3==1.26.62
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.2
[pip3] torch==2.0.0
[pip3] torch-tb-profiler==0.4.1
[pip3] torchdata==0.6.0
[pip3] torchmetrics==0.11.1
[pip3] torchsnapshot-nightly==2023.2.17
[pip3] torchx==0.5.0.dev0
[pip3] triton==2.0.0
[conda] No relevant packages



cc @ezyang @gchanan @zou3519 @SsnL @VitalyFedyunin @ejguan @NivekT @dzhulgakov


Discussion:
cc) @wenleix 
@kiukchung To unblock you ASAP, would you pls subclass from `_ShardingIterDataPipe`? I plan to make `apply_sharding` running with `DataPipe` that with `apply_sharding` method.

And, yes. We should include it TorchData's release note for BC breaking. cc: @NivekT 
Thanks @ejguan! That's exactly what I've done to get myself unblocked for now. But I found a few more interesting quirks when I subclassed `_ShardingIterDataPipe` that might be worth documenting in the instructions to "migrate" over to the new shardable paradigm:

1. The `is_shardable()` method is now gone (things will work if you forget to remove it, but for cleanliness you should remove it).
2. The `apply_sharding()` method introduces a new `sharding_group` parameter to the method signature. At first glance, the method signature is BC since `sharding_group` defaults to `SHARDING_PRIORITIES.DEFAULT`. But the way `DataLoader` calls `apply_sharding()` breaks BC if you don't modify the `apply_sharding()` logic. Prior to 2.0, the DataLoader would call `apply_sharding()` once if `is_shardable` attribute existed and returned `True`. And the `num_instances` and `instance_id` would already account for process group rank and world_size. In 2.0 the dataloader calls `apply_sharding()` twice and what gets passed as the `instance_id` and `num_workers` parameter is dependent on the `sharding_group`:
        a. IF `sharding_group=MULTIPROCESSING` then instance_id and num_workers is the dataloader worker's rank and num_workers (on the host).
        b. IF `sharding_group=DISTRIBUTED` then instance_id and num_workers is the process group's (trainer's) rank and world size

So in essence if my prior implementation looked like this:

```
def apply_sharding(self, num_instances, instance_id):
      self.global_world_size = num_instances # world_size * num_dataloader_workers
      self.global_rank = instance_id # dataloader_worker_rank * num_dataloader_workers + rank * world_size

# sharding would be done by 
shard_idx % self.global_world_size == self.global_rank
```

It would NOT be enough to simply add the `sharding_group` param, but rather it has to change to account for the two calls that the dataloader worker will make:

```
def apply_sharding(self, num_instances, instance_id, sharding_group=SHARDING_PRIORITY.DEFAULT):
     if sharding_group == SHARDING_PRIORITY.MULTIPROCESSING:
             self.instance_id = instance_id if self.instance_id
              self.num_instances = num_instances
     elif sharding_group == SHARDING_PRIORITY.DISTRIBUTED:
             self.rank = rank
             self.world_size = world_size
     else: # SHARDING_PRIORITY.DEFAULT
             self.global_world_size = num_instances
             self.global_rank = instance_id
 
# sharding would be done by
if self.global_world_size == -1:
     self.global_world_size = self.num_instances * self.world_size
if self.global_rank == -1:
     self.global_rank = self.instance_id * self.num_instances + self.rank * self.world_size
shard_idx % self.global_world_size == self.global_rank
```
@kiukchung Sorry for the inconvenience, we should realize it when the change is made in https://github.com/pytorch/pytorch/pull/88424

Since we have already made this BC-breaking change, to resolve this issue going forward, we will do the follow work:
1. Update the release note from TorchData to reflect those changes
2. Update `apply_sharding` to accept any DataPipe with `apply_sharding` method tracked in https://github.com/pytorch/data/issues/1081
    - Add a feature to inspect the `dp.apply_sharding` method. IF the last `sharding_group` is not acceptable, we will treat it as `DEFAULT` and skip sending this extra argument to the `dp.apply_sharding`
BTW, are you still relying on `is_shardable` to determine if you want to apply sharding to the DataPipe?
@ejguan thanks for the quick updates! sounds like a plan for me. I'm currently unblocked since I went ahead and subclassed from `_ShardingIterDataPipe`. 

> are you still relying on is_shardable to determine if you want to apply sharding to the DataPipe?

Not anymore, I removed the `is_shardable()` method since it is no longer used by the DataLoader. I don't manually call the `apply_sharding()` method anywhere.

I used to have a runtime_checkable Protocol called `Shardable` (as shown at the top) but this is also obsolete since I can always do an `isinstance(_ShardingIterDataPipe, pipe)` if needed (but I don't have a use-case for this other than in certain sanity tests).

I'll comment on https://github.com/pytorch/data/issues/1081 for specifics on the forward fix.

Closing in favor of tracking https://github.com/pytorch/data/issues/1081 since I have a workaround for now.



I am actually thinking about forward fixing it directly since PyTorch might do a minor release.
@ejguan @kiukchung   Please note: 

This issue is in the  milestones :[ https://github.com/pytorch/pytorch/milestone/36?closed=1](https://github.com/pytorch/pytorch/milestone/36?closed=1), if you want to see your fix included in this minor release. Please post it as a cherry-pick into the [[v2.0.1] Release Tracker](https://github.com/pytorch/pytorch/issues/97272). 

**The deadline is April 14, 5PM PST.**

**Only issues that have ‘cherry-picks’ will be considered for the release.**

Common FAQs:

Q1: Where can I find more information on the release process and terminology?

A: [pytorch/RELEASE.md at master · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/blob/master/RELEASE.md)

Q2: Am I guaranteed to be included in the cherry-pick if I do above?

A: No, it is not guaranteed, the Release Team will review all submissions against the listed criteria before making the final decision on what to include on 4/17.

Q3: When is 2.1 going to be released?

A: We do not have a formal date at this time but will update the community when we do. Our immediate focus is 2.0.1. Note that 1.12 was released on 6/28/22, 1.13 on 10/28/22 and 2.0 on 3/15/23.

Q4: **I missed the 4/14 5PM PST deadline, is there any option to have an extension?**

A: **No, in order to meet our 4/28 goal, we must hold 4/14 as our deadline and will not accept any requests after the fact. We are over communicating the timelines and process with the community to avoid such issues.**

Q5: Where should I double check to see if my issue is in the cherry pick tracker?

A: [[v2.0.1] Release Tracker · Issue #97272 · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/issues/97272)

Q6: Where can I find the Release Compatibility Matrix for PyTorch?

A: [pytorch/RELEASE.md at master · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/blob/master/RELEASE.md#release-compatibility-matrix)

Please contact OSS Releng team members if you have any questions/comments. Again we appreciate everyone’s time and commitment to the community, PyTorch and 2.0 and 2.01 releases!


Please refer to this post for more details: https://dev-discuss.pytorch.org/t/pytorch-release-2-0-1-important-information/1176

Resolving patch:
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
--- a/torch/utils/data/dataloader.py
+++ b/torch/utils/data/dataloader.py
@@ -34,7 +34,6 @@
     Dataset,)
 
 from torch.utils.data.datapipes.datapipe import _IterDataPipeSerializationWrapper, _MapDataPipeSerializationWrapper
-from torch.utils.data.datapipes.iter.sharding import SHARDING_PRIORITIES
 
 from . import _utils
 
@@ -104,6 +103,7 @@ def _get_distributed_settings():
 
 
 def _sharding_worker_init_fn(worker_init_fn, world_size, rank_id, worker_id):
+    global_worker_id = worker_id
     info = torch.utils.data.get_worker_info()
     assert info is not None
     total_workers = info.num_workers
@@ -111,10 +111,10 @@ def _sharding_worker_init_fn(worker_init_fn, world_size, rank_id, worker_id):
     assert isinstance(datapipe, (IterDataPipe, MapDataPipe))
     # To distribute elements across distributed process evenly, we should shard data on distributed
     # processes first then shard on worker processes
-    torch.utils.data.graph_settings.apply_sharding(
-        datapipe, world_size, rank_id, sharding_group=SHARDING_PRIORITIES.DISTRIBUTED)
-    torch.utils.data.graph_settings.apply_sharding(
-        datapipe, total_workers, worker_id, sharding_group=SHARDING_PRIORITIES.MULTIPROCESSING)
+    total_workers *= world_size
+    global_worker_id = global_worker_id * world_size + rank_id
+    # For BC, use default SHARDING_PRIORITIES
+    torch.utils.data.graph_settings.apply_sharding(datapipe, total_workers, global_worker_id)
     if worker_init_fn is not None:
         worker_init_fn(worker_id)
 
@@ -666,8 +666,8 @@ def __init__(self, loader):
         # Adds forward compatibilities so classic DataLoader can work with DataPipes:
         #   Taking care of distributed sharding
         if isinstance(self._dataset, (IterDataPipe, MapDataPipe)):
-            torch.utils.data.graph_settings.apply_sharding(
-                self._dataset, self._world_size, self._rank, sharding_group=SHARDING_PRIORITIES.DISTRIBUTED)
+            # For BC, use default SHARDING_PRIORITIES
+            torch.utils.data.graph_settings.apply_sharding(self._dataset, self._world_size, self._rank)
 
         self._dataset_fetcher = _DatasetKind.create_fetcher(
             self._dataset_kind, self._dataset, self._auto_collation, self._collate_fn, self._drop_last)
diff --git a/torch/utils/data/datapipes/iter/sharding.py b/torch/utils/data/datapipes/iter/sharding.py
--- a/torch/utils/data/datapipes/iter/sharding.py
+++ b/torch/utils/data/datapipes/iter/sharding.py
@@ -13,15 +13,18 @@
     "ShardingFilterIterDataPipe",
 ]
 
+
 class SHARDING_PRIORITIES(IntEnum):
     DEFAULT = 1
     DISTRIBUTED = 2
     MULTIPROCESSING = 3
 
+
 class _ShardingIterDataPipe(IterDataPipe):
-    def apply_sharding(self, num_of_instances, instance_id, sharding_group):
+    def apply_sharding(self, num_of_instances: int, instance_id: int, sharding_group: SHARDING_PRIORITIES):
         raise NotImplementedError
 
+
 @functional_datapipe('sharding_filter')
 class ShardingFilterIterDataPipe(_ShardingIterDataPipe):
     r"""
diff --git a/torch/utils/data/graph_settings.py b/torch/utils/data/graph_settings.py
--- a/torch/utils/data/graph_settings.py
+++ b/torch/utils/data/graph_settings.py
@@ -35,6 +35,14 @@ def _get_all_graph_pipes_helper(graph: DataPipeGraph, id_cache: Set[int]) -> Lis
     return results
 
 
+def _is_sharding_datapipe(datapipe: DataPipe) -> bool:
+    if isinstance(datapipe, _ShardingIterDataPipe):
+        return True
+    if hasattr(datapipe, "apply_sharding") and inspect.ismethod(datapipe.apply_sharding):
+        return True
+    return False
+
+
 def apply_sharding(datapipe: DataPipe,
                    num_of_instances: int,
                    instance_id: int,
@@ -48,11 +56,16 @@ def apply_sharding(datapipe: DataPipe,
     def _helper(graph, prev_applied=None):
         for _, (dp, sub_graph) in graph.items():
             applied = None
-            if isinstance(dp, _ShardingIterDataPipe):
+            if _is_sharding_datapipe(dp):
                 if prev_applied is not None:
                     raise RuntimeError("Sharding twice on a single pipeline is likely unintended and will cause data loss. "
                                        f"Sharding already applied to {prev_applied} while trying to apply to {dp}")
-                dp.apply_sharding(num_of_instances, instance_id, sharding_group=sharding_group)
+                # For BC, only provide sharding_group if accepted
+                sig = inspect.signature(dp.apply_sharding)
+                if len(sig.parameters) < 3:
+                    dp.apply_sharding(num_of_instances, instance_id)
+                else:
+                    dp.apply_sharding(num_of_instances, instance_id, sharding_group=sharding_group)
                 applied = dp
             if applied is None:
                 applied = prev_applied


Test patch:
diff --git a/test/test_datapipe.py b/test/test_datapipe.py
--- a/test/test_datapipe.py
+++ b/test/test_datapipe.py
@@ -2656,6 +2656,23 @@ def test_circular_serialization_with_dill(self):
         })}
         self.assertEqual(res2, exp_res_2)
 
+
+class CustomShardingIterDataPipe(IterDataPipe):
+    def __init__(self, dp):
+        self.dp = dp
+        self.num_of_instances = 1
+        self.instance_id = 0
+
+    def apply_sharding(self, num_of_instances, instance_id):
+        self.num_of_instances = num_of_instances
+        self.instance_id = instance_id
+
+    def __iter__(self):
+        for i, d in enumerate(self.dp):
+            if i % self.num_of_instances == self.instance_id:
+                yield d
+
+
 class TestSharding(TestCase):
 
     def _get_pipeline(self):
@@ -2762,6 +2779,12 @@ def construct_sharded_pipe():
         with self.assertRaises(Exception):
             dp.apply_sharding(2, 1, sharding_group=LEGACY_SHARDING_PRIORITIES.DEFAULT)
 
+    def test_legacy_custom_sharding(self):
+        dp = self._get_pipeline()
+        sharded_dp = CustomShardingIterDataPipe(dp)
+        torch.utils.data.graph_settings.apply_sharding(sharded_dp, 3, 1)
+        items = list(sharded_dp)
+        self.assertEqual([1, 20], items)
 
     def test_sharding_length(self):
         numbers_dp = dp.iter.IterableWrapper(range(13))
@@ -2796,6 +2819,19 @@ def test_old_dataloader(self):
 
         self.assertEqual(sorted(expected), sorted(items))
 
+    def test_legacy_custom_sharding_with_old_dataloader(self):
+        dp0 = self._get_pipeline()
+        expected = list(dp0)
+
+        dp0 = self._get_pipeline()
+        dp0 = CustomShardingIterDataPipe(dp0)
+        dl = DataLoader(dp0, batch_size=1, shuffle=False, num_workers=2)
+        items = []
+        for i in dl:
+            items.append(i)
+
+        self.assertEqual(sorted(expected), sorted(items))
+
     def test_multi_sharding(self):
         # Raises Error when multiple sharding on the single branch
         numbers_dp = dp.iter.IterableWrapper(range(13))