Restrict Transformers to data-parallel replicas

configuration_limite.py

@@ -129,18 +129,6 @@ class LimiteConfig(PretrainedConfig):
     model_type = MODEL_TYPE
     keys_to_ignore_at_inference = ["past_key_values"]
 
-    base_model_tp_plan = {
-        "layers.*.self_attn.qkv_proj": "colwise",
-        "layers.*.self_attn.o_proj": "rowwise",
-        "layers.*.mlp.gate_up_proj": "packed_colwise",
-        "layers.*.mlp.down_proj": "rowwise",
-    }
-    base_model_pp_plan = {
-        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
-        "layers": (["hidden_states"], ["hidden_states"]),
-        "norm": (["hidden_states"], ["hidden_states"]),
-    }
-
     @classmethod
     def from_dict(cls, config_dict: dict[str, Any], **kwargs: Any) -> "LimiteConfig":
         missing = [

modeling_limite.py

@@ -671,6 +671,46 @@ class LimitePreTrainedModel(PreTrainedModel):
     _supports_attention_backend = True
     _can_compile_fullgraph = True
 
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | None,
+        *model_args: Any,
+        **kwargs: Any,
+    ) -> LimitePreTrainedModel:
+        requested_parallelism = [
+            name
+            for name in ("tp_plan", "tp_size", "distributed_config")
+            if kwargs.get(name) is not None
+        ]
+        device_map = kwargs.get("device_map")
+        if isinstance(device_map, str) and device_map in {
+            "auto",
+            "balanced",
+            "balanced_low_0",
+            "sequential",
+        }:
+            requested_parallelism.append(f"device_map={device_map!r}")
+        elif isinstance(device_map, dict):
+            placements = {str(device) for device in device_map.values()}
+            if len(placements) > 1:
+                requested_parallelism.append("multi-device device_map")
+
+        if requested_parallelism:
+            requested = ", ".join(requested_parallelism)
+            raise NotImplementedError(
+                "Limite supports one complete model replica per process; "
+                "tensor parallelism, pipeline parallelism, and multi-device "
+                f"model sharding are not supported (requested: {requested}). "
+                "Use process-level data parallelism with one explicit device "
+                "per replica."
+            )
+        return super().from_pretrained(
+            pretrained_model_name_or_path,
+            *model_args,
+            **kwargs,
+        )
+
     @torch.no_grad()
     def _init_weights(self, module: nn.Module) -> None:
         super()._init_weights(module)