Style

Author: remi-or

examples/pytorch/continuous_batching.py

@@ -55,15 +55,18 @@ def setup_metrics():
         from opentelemetry.sdk.trace import TracerProvider
         from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
-
         resource = Resource.create({"service.name": "transformers"})
         metrics_exporter = PeriodicExportingMetricReader(
-            OTLPMetricExporter(endpoint="http://localhost:9090/api/v1/otlp/v1/metrics"),  # Uses OTEL_EXPORTER_OTLP_METRICS_ENDPOINT env var
-            export_interval_millis=1000
+            OTLPMetricExporter(
+                endpoint="http://localhost:9090/api/v1/otlp/v1/metrics"
+            ),  # Uses OTEL_EXPORTER_OTLP_METRICS_ENDPOINT env var
+            export_interval_millis=1000,
         )
         meter_provider = MeterProvider(resource=resource, metric_readers=[metrics_exporter])
         metrics.set_meter_provider(meter_provider)
-        trace_exporter = OTLPSpanExporter(endpoint="http://localhost:4318/v1/traces")  # Uses OTEL_EXPORTER_OTLP_TRACES_ENDPOINT env var
+        trace_exporter = OTLPSpanExporter(
+            endpoint="http://localhost:4318/v1/traces"
+        )  # Uses OTEL_EXPORTER_OTLP_TRACES_ENDPOINT env var
         tracer_provider = TracerProvider(resource=resource)
         tracer_provider.add_span_processor(BatchSpanProcessor(trace_exporter))
         trace.set_tracer_provider(tracer_provider)
@@ -213,9 +216,7 @@ def batch_generate(
     # If no output file is provided, we pick a name based on the args
     if args.output_file is None:
         os.makedirs("runs/cb", exist_ok=True)
-        args.output_file = (
-            f"runs/cb/{args.num_blocks}_{args.max_batch_tokens}_{args.attn}_{args.matmul_precision}_{args.samples}.json"
-        )
+        args.output_file = f"runs/cb/{args.num_blocks}_{args.max_batch_tokens}_{args.attn}_{args.matmul_precision}_{args.samples}.json"
 
     # Run warmup batch generation
     batch_generate(

src/transformers/generation/continuous_batching/cache.py

@@ -274,7 +274,9 @@ def compute_num_blocks_and_max_batch_tokens(
         logger.info(f"Cache memory: {cache_memory}")
 
         # Compute memory footprints # TODO: check and explain better
-        mem_per_activation_token = self._activation_dtype.itemsize * (self.hidden_size + self.vocab_size) * self._activation_safety_factor
+        mem_per_activation_token = (
+            self._activation_dtype.itemsize * (self.hidden_size + self.vocab_size) * self._activation_safety_factor
+        )
         mem_per_cache_token = 2 * self.num_heads * self.head_dim * self.num_layers * cache_dtype.itemsize
         mem_per_input_token = 8 * m * self._input_dtype.itemsize
         logger.info(f"Memory per activation token: {mem_per_activation_token}")