Fix CI

@@ -92,6 +92,7 @@     get_fsdp2_grad_scaler,     get_grad_scaler,     get_mixed_precision_context_manager,+    get_model_tp_size,     get_pretty_name,     has_offloaded_params,     is_bf16_available,@@ -1900,14 +1901,15 @@ def prepare_model(                     if self.ddp_handler is not None:                         self.ddp_handler.register_comm_hook(model)             elif self.parallelism_config and self.parallelism_config.tp_enabled:-                if not hasattr(model, "tp_size"):+                model_tp_size = get_model_tp_size(model)+                if model_tp_size is None:                     raise NotImplementedError(                         "Model should undergo tensor parallel before passing it to accelerate."                         "You can use .from_pretrained(..., tp_plan='auto') if the model supports"                     )-                if model.tp_size != self.parallelism_config.tp_size:+                if model_tp_size != self.parallelism_config.tp_size:                     raise ValueError(-                        f"tp_size in the plugin {self.parallelism_config.tp_size} should be same as model's tp size {model.tp_size}"+                        f"tp_size in the plugin {self.parallelism_config.tp_size} should be same as model's tp size {model_tp_size}"                     )             elif self.is_fsdp2:                 raise ValueError(@@ -2338,7 +2340,9 @@ def _prepare_deepspeed(self, *args):                     {                         "scheduler.params.warmup_min_lr": 0,                         "scheduler.params.warmup_max_lr": max_lr,-                        "scheduler.params.warmup_num_steps": scheduler.warmup_num_steps,+                        # `DummyScheduler` defaults to 0, which `deepspeed>=0.19.6` rejects outright.+                        # Older versions silently clamped it to 2, so keep the value positive.+                        "scheduler.params.warmup_num_steps": max(1, scheduler.warmup_num_steps),                     }                 )                 if scheduler.total_num_steps is not None:

@@ -283,6 +283,7 @@     compile_regions_fsdp2,     convert_bytes,     extract_model_from_parallel,+    get_model_tp_size,     get_module_children_bottom_up,     get_pretty_name,     has_compiled_regions,

@@ -248,6 +248,29 @@ def model_has_dtensor(model: torch.nn.Module) -> bool:     return any(isinstance(p, DTensor) for p in model.parameters())  +def get_model_tp_size(model: torch.nn.Module) -> Optional[int]:+    """+    Get the tensor parallel degree a `transformers` model was sharded with, or `None` if it was not sharded.++    Args:+        model (`torch.nn.Module`):+            The model to inspect.++    Returns:+        `Optional[int]`: The model's tensor parallel size.+    """+    # `transformers<5` records it on the model itself, while `transformers>=5` moved it to the+    # `DistributedConfig` held by the model config and left `model.tp_size` behind as a `None` stub.+    tp_size = getattr(model, "tp_size", None)+    if tp_size is not None:+        return tp_size++    distributed_config = getattr(getattr(model, "config", None), "distributed_config", None)+    if isinstance(distributed_config, dict):+        return distributed_config.get("tp_size")+    return getattr(distributed_config, "tp_size", None)++ def extract_model_from_parallel(     model, keep_fp32_wrapper: bool = True, keep_torch_compile: bool = True, recursive: bool = False ):

@@ -208,7 +208,7 @@ def test_deepspeed_plugin(self, stage):             "optimizer.params.weight_decay": 0.0,             "scheduler.params.warmup_min_lr": 0.0,             "scheduler.params.warmup_max_lr": 5e-5,-            "scheduler.params.warmup_num_steps": 0,+            "scheduler.params.warmup_num_steps": 1,             "train_micro_batch_size_per_gpu": 16,             "gradient_clipping": 1.0,             "train_batch_size": 16,@@ -288,7 +288,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):             "optimizer.params.weight_decay": 0.0,             "scheduler.params.warmup_min_lr": 0.0,             "scheduler.params.warmup_max_lr": 5e-5,-            "scheduler.params.warmup_num_steps": 0,+            "scheduler.params.warmup_num_steps": 1,             "train_micro_batch_size_per_gpu": 16,             "gradient_clipping": 1.0,             "train_batch_size": 16,@@ -567,7 +567,7 @@ def test_save_checkpoints(self):             "optimizer.params.weight_decay": 0.0,             "scheduler.params.warmup_min_lr": 0.0,             "scheduler.params.warmup_max_lr": 5e-5,-            "scheduler.params.warmup_num_steps": 0,+            "scheduler.params.warmup_num_steps": 1,             "train_micro_batch_size_per_gpu": 16,             "gradient_clipping": 1.0,             "train_batch_size": 16,

@@ -17,6 +17,7 @@ import unittest import warnings from collections import UserDict, namedtuple+from types import SimpleNamespace from typing import NamedTuple, Optional from unittest.mock import Mock, patch @@ -48,6 +49,7 @@     convert_to_fp32,     extract_model_from_parallel,     find_device,+    get_model_tp_size,     has_offloaded_params,     is_torch_xla_available,     listify,@@ -443,6 +445,23 @@ def test_has_offloaded_params(self):         attach_align_device_hook(model, offload=True)         assert has_offloaded_params(model) +    def test_get_model_tp_size(self):+        model = RegressionModel()+        assert get_model_tp_size(model) is None++        # `transformers<5` records the degree on the model itself+        model.tp_size = 2+        assert get_model_tp_size(model) == 2++        # `transformers>=5` leaves `model.tp_size` behind as a `None` stub and moves the degree to the config+        model.tp_size = None+        model.config = SimpleNamespace(distributed_config=SimpleNamespace(tp_size=4))+        assert get_model_tp_size(model) == 4++        # a config that round-tripped through JSON holds a plain dict+        model.config = SimpleNamespace(distributed_config={"tp_size": 8})+        assert get_model_tp_size(model) == 8+     def test_concatenate(self):         tensor1 = torch.randn(2, 3)         tensor2 = torch.randn(2, 3)