Repository: huggingface/transformers
Pull request: enable tp for benchmark

enable tp in benchmark_v2, to ensure large model could run.

Changed files:
filename: benchmark_v2/framework/benchmark_config.py
patch: @@ -69,6 +69,7 @@ def __init__(
         attn_implementation: str = "eager",
         compile_kwargs: dict[str, Any] | None = None,
         kernelize: bool = False,
+        tp_plan: str | dict[str, str] | None = None,
         name: str | None = None,
         skip_validity_check: bool = False,
     ) -> None:
@@ -83,6 +84,7 @@ def __init__(
         self.num_tokens_to_generate = num_tokens_to_generate
         # Generation parameters
         self.attn_implementation = attn_implementation
+        self.tp_plan = tp_plan
         # Optimization parameters
         if compile_kwargs is None:
             self.compile_config = None
@@ -100,7 +102,6 @@ def __init__(
     def check_validity(self, skip_validity_check: bool = False) -> None:
         if skip_validity_check:
             return
-
         # If flash_attention_2 is selected but not available, default to SDPA
         if self.attn_implementation == "flash_attention_2" and not is_fa2_or_kernel_available():
             logger.error("Flash attention is not available. Defaulting to SDPA.")
@@ -150,6 +151,7 @@ def infer_name(self, compact: bool = True) -> str:
             compile_str = f"compiled_{self.compile_config.mode}" if self.compile_config is not None else "uncompiled"
             kernelize_str = "kernelized" if self.kernelize else "unkernelized"
             continuous_batching_str = "cb" if self.continuous_batching else "generate"
+            tp_str = "tp" if self.tp_plan is not None else "no_tp"
             sep = "-"
         else:
             iter_str = f"{self.warmup_iterations} warmup, {self.measurement_iterations} iterations"
@@ -159,9 +161,22 @@ def infer_name(self, compact: bool = True) -> str:
             compile_str = "compiled" if self.compile_config is not None else "not compiled"
             kernelize_str = "kernelized" if self.kernelize else "not kernelized"
             continuous_batching_str = "continuous batching" if self.continuous_batching else "regular generate"
+            if self.tp_plan is None:
+                tp_str = "no_tp"
+            else:
+                tp_str = "tp_custom" if isinstance(self.tp_plan, dict) else "tp_auto"
             sep = ", "
         return sep.join(
-            [iter_str, gpu_monitor_str, dimensions_str, attn_code, compile_str, kernelize_str, continuous_batching_str]
+            [
+                iter_str,
+                gpu_monitor_str,
+                dimensions_str,
+                attn_code,
+                compile_str,
+                kernelize_str,
+                continuous_batching_str,
+                tp_str,
+            ]
         )
 
     def to_dict(self) -> dict[str, Any]:
@@ -177,6 +192,7 @@ def to_dict(self) -> dict[str, Any]:
             "attn_implementation": self.attn_implementation,
             "compile_kwargs": self.compile_config.to_dict() if self.compile_config is not None else None,
             "kernelize": self.kernelize,
+            "tp_plan": self.tp_plan,
         }
 
     @classmethod
@@ -192,6 +208,7 @@ def from_dict(cls, data: dict[str, Any], skip_validity_check: bool = False) -> "
             attn_implementation=data.get("attn_implementation", "eager"),
             compile_kwargs=data.get("compile_kwargs"),
             kernelize=data.get("kernelize", False),
+            tp_plan=data.get("tp_plan"),
             name=data.get("name"),
             skip_validity_check=skip_validity_check,
         )

filename: benchmark_v2/framework/benchmark_runner.py
patch: @@ -169,6 +169,12 @@ def cleanup(self) -> None:
         self.model = None
         flush_memory()
 
+    @staticmethod
+    def _is_primary_process() -> bool:
+        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+            return True
+        return torch.distributed.get_rank() == 0
+
     def setup_benchmark(self, model_id: str, config: BenchmarkConfig) -> None:
         # Some attributes only need to be set once per model
         if self._setup_for != model_id:
@@ -184,7 +190,7 @@ def setup_benchmark(self, model_id: str, config: BenchmarkConfig) -> None:
             max_length=config.sequence_length,
             truncation=True,
             return_attention_mask=True,
-        ).to(config.device)
+        )
         self.inputs["use_cache"] = True
 
         # Prepare generation config
@@ -206,15 +212,18 @@ def setup_benchmark(self, model_id: str, config: BenchmarkConfig) -> None:
         self.logger.debug(f"Loading model {model_id} on device {config.device}...")
         dtype = getattr(torch, config.dtype.removeprefix("torch."))
         use_kernels = config.kernelize and kernelize is not None and Mode is not None
+        device_map = config.device if config.tp_plan is None else None
         self.model = AutoModelForCausalLM.from_pretrained(
             model_id,
             dtype=dtype,
             attn_implementation=config.attn_implementation,
             generation_config=generation_config,
             use_kernels=use_kernels,
-            device_map=config.device,
+            device_map=device_map,
+            tp_plan=config.tp_plan,
         )
         self.model = self.model.eval()
+        self.inputs = self.inputs.to(self.model.device)
 
     def run_benchmark(self, config: BenchmarkConfig, num_tokens_to_profile: int = 0) -> BenchmarkResult | None:
         """Run a single benchmark with the given model ID and config."""
@@ -378,6 +387,8 @@ def run_benchmarks(
             raise RuntimeError("No benchmark was run successfully")
 
         if pretty_print_summary:
+            if not self._is_primary_process():
+                return (timestamp, all_results)
             print()
             print("=" * 100)
             print(f"Finished benchmarks in {time.perf_counter() - start_time:.2f} seconds")
@@ -401,6 +412,8 @@ def run_benchmarks(
 
     def save_results(self, model_name: str, results: dict, timestamp: str = "", summarized: bool = True) -> str:
         """Save benchmark results to JSON file."""
+        if not self._is_primary_process():
+            return ""
         # Create model-specific subdirectory
         model_name = model_name.replace("/", "_")
         model_dir = os.path.join(self.output_dir, model_name)

filename: benchmark_v2/run_benchmarks.py
patch: @@ -51,6 +51,7 @@
     )
     parser.add_argument("--config-file", type=str, help="Path to a config file stored as a json or jsonl format")
     parser.add_argument("--num-tokens-to-profile", "-p", type=int, default=0, help="Number of tokens to profile")
+    parser.add_argument("--enable-tp", action="store_true", help="Enable tensor parallelism with tp_plan=auto")
 
     parser.add_argument("--branch-name", type=str, help="Git branch name")
     parser.add_argument("--commit-id", type=str, help="Git commit ID (if not provided, will auto-detect from git)")
@@ -118,11 +119,15 @@
         not args.no_gpu_monitoring,
     )
 
+    if args.enable_tp:
+        for config in configs:
+            config.tp_plan = "auto"
+
     runner = BenchmarkRunner(logger, args.output_dir, args.branch_name, args.commit_id, args.commit_message)
     timestamp, results = runner.run_benchmarks(
         args.model_id, configs, args.num_tokens_to_profile, pretty_print_summary=True
     )
 
     dataset_id = args.push_result_to_dataset
-    if dataset_id is not None and len(results) > 0:
+    if dataset_id is not None and len(results) > 0 and runner._is_primary_process():
         runner.push_results_to_hub(dataset_id, results, timestamp)