Repository: huggingface/transformers
Pull request: Fix fsdp for generic-task models

# What does this PR do?

As per the title. Fixes https://github.com/huggingface/transformers/issues/39977
Inheriting from `ABC` would result in this kind of errors when using fsdp: `TypeError: __class__ assignment: 'FSDPQwen2ForTokenClassification' object layout differs from 'Qwen2ForTokenClassification'`.

As having them inherit from `ABC` was sugar to indicate they should only be used through proper inheritance, it's safe to remove it

Changed files:
filename: src/transformers/modeling_layers.py
patch: @@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from abc import ABC
 from functools import partial
 from typing import Optional
 
@@ -96,7 +95,7 @@ def __call__(self, *args, **kwargs):
 
 
 @auto_docstring
-class GenericForSequenceClassification(ABC):
+class GenericForSequenceClassification(object):
     base_model_prefix = "model"
 
     def __init__(self, config):
@@ -171,7 +170,7 @@ def forward(
 
 
 @auto_docstring
-class GenericForQuestionAnswering(ABC):
+class GenericForQuestionAnswering(object):
     base_model_prefix = "model"
 
     def __init__(self, config):
@@ -232,7 +231,7 @@ def forward(
 
 
 @auto_docstring
-class GenericForTokenClassification(ABC):
+class GenericForTokenClassification(object):
     base_model_prefix = "model"
 
     def __init__(self, config):

filename: src/transformers/testing_utils.py
patch: @@ -3471,3 +3471,23 @@ def find_expectation(self, properties: DeviceProperties = (None, None, None)) ->
 
     def __repr__(self):
         return f"{self.data}"
+
+
+def torchrun(script: str, nproc_per_node: int, is_torchrun: bool = True, env: Optional[dict] = None):
+    """Run the `script` using `torchrun` command for multi-processing in a subprocess. Captures errors as necessary."""
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".py") as tmp:
+        tmp.write(script)
+        tmp.flush()
+        tmp.seek(0)
+        if is_torchrun:
+            cmd = (
+                f"torchrun --nproc_per_node {nproc_per_node} --master_port {get_torch_dist_unique_port()} {tmp.name}"
+            ).split()
+        else:
+            cmd = ["python3", tmp.name]
+
+        # Note that the subprocess will be waited for here, and raise an error if not successful
+        try:
+            _ = subprocess.run(cmd, capture_output=True, env=env, text=True, check=True)
+        except subprocess.CalledProcessError as e:
+            raise Exception(f"The following error was captured: {e.stderr}")

filename: tests/generation/test_fsdp.py
patch: @@ -13,6 +13,7 @@
 # limitations under the License.
 
 import argparse
+import textwrap
 from typing import Any, Callable
 
 from transformers import is_torch_available, is_torch_xpu_available
@@ -24,6 +25,7 @@
     get_torch_dist_unique_port,
     require_torch_multi_accelerator,
     torch_device,
+    torchrun,
 )
 from transformers.utils import is_ccl_available, is_ipex_available
 
@@ -141,6 +143,33 @@ def test_fsdp2_generate(self):
         # successful return here == success - any errors would have caused an error in the sub-call
 
 
+class TestFSDPGenericTaskModel(TestCasePlus):
+    nproc_per_node = 2
+
+    def test_generic_task_model_can_be_sharded(self):
+        script_to_run = textwrap.dedent(
+            """
+            import torch
+            from torch.distributed.fsdp import fully_shard
+            from transformers import AutoModelForTokenClassification
+
+            torch.distributed.init_process_group(
+                backend="nccl" if torch.cuda.is_available() else "gloo", init_method="env://"
+            )
+            rank = torch.distributed.get_rank()
+            if torch.cuda.is_available():
+                torch.cuda.set_device(rank)
+
+            # Make sure it works
+            model = AutoModelForTokenClassification.from_pretrained("Qwen/Qwen2-0.5B")
+            module = fully_shard(model)
+
+            torch.distributed.destroy_process_group()
+            """
+        )
+        torchrun(script_to_run, self.nproc_per_node, env=self.get_env())
+
+
 if __name__ == "__main__":
     # The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
     #

filename: tests/tensor_parallel/test_tensor_parallel.py
patch: @@ -15,7 +15,6 @@
 # Run the test: CUDA_VISIBLE_DEVICES=0,1 RUN_SLOW=1 pytest -sv tests/tensor_parallel/test_tensor_parallel.py
 
 import os
-import subprocess
 import tempfile
 import textwrap
 
@@ -24,10 +23,10 @@
 from transformers.testing_utils import (
     TestCasePlus,
     backend_device_count,
-    get_torch_dist_unique_port,
     require_huggingface_hub_greater_or_equal,
     require_torch_multi_accelerator,
     torch_device,
+    torchrun,
 )
 
 
@@ -67,25 +66,6 @@ def size(self):
 class TestTensorParallel(TestCasePlus):
     nproc_per_node = 2
 
-    def torchrun(self, script: str, is_torchrun: bool = True):
-        """Run the `script` using `torchrun` command for multi-processing in a subprocess. Captures errors as necessary."""
-        with tempfile.NamedTemporaryFile(mode="w+", suffix=".py") as tmp:
-            tmp.write(script)
-            tmp.flush()
-            tmp.seek(0)
-            if is_torchrun:
-                cmd = (
-                    f"torchrun --nproc_per_node {self.nproc_per_node} --master_port {get_torch_dist_unique_port()} {tmp.name}"
-                ).split()
-            else:
-                cmd = ["python3", tmp.name]
-
-            # Note that the subprocess will be waited for here, and raise an error if not successful
-            try:
-                _ = subprocess.run(cmd, capture_output=True, env=self.get_env(), text=True, check=True)
-            except subprocess.CalledProcessError as e:
-                raise Exception(f"The following error was captured: {e.stderr}")
-
     def test_model_forward(self):
         script_to_run = textwrap.dedent(
             """
@@ -124,7 +104,7 @@ def test_model_forward(self):
             torch.distributed.destroy_process_group()
             """
         )
-        self.torchrun(script_to_run)
+        torchrun(script_to_run, self.nproc_per_node, env=self.get_env())
 
     def test_model_backward_pass(self):
         script_to_run = textwrap.dedent(
@@ -150,7 +130,7 @@ def test_model_backward_pass(self):
             torch.distributed.destroy_process_group()
             """
         )
-        self.torchrun(script_to_run)
+        torchrun(script_to_run, self.nproc_per_node, env=self.get_env())
 
     def test_model_generate(self):
         script_to_run = textwrap.dedent(
@@ -190,7 +170,7 @@ def test_model_generate(self):
             torch.distributed.destroy_process_group()
             """
         )
-        self.torchrun(script_to_run)
+        torchrun(script_to_run, self.nproc_per_node, env=self.get_env())
 
     @require_huggingface_hub_greater_or_equal("0.31.4")
     def test_model_save(self):
@@ -217,7 +197,7 @@ def test_model_save(self):
                     model.save_pretrained(result_dir)
                     """
                 )
-                self.torchrun(script_to_run, is_torchrun=is_torchrun)
+                torchrun(script_to_run, self.nproc_per_node, is_torchrun=is_torchrun, env=self.get_env())
 
             non_tp_model_path = os.path.join(tmp_dir, "nontp")
             tp_model_path = os.path.join(tmp_dir, "tp")