Apply black

Author: FKKimura

onecomp/analyzer/quantization_error.py

@@ -107,15 +107,21 @@ def _plot_layer_type(ax, layer_type, label_points, colors):
                 cy = min(y, ylim * 0.92)
                 text = f"{k} ({x:.1f}, {y:.2f})"
                 ax.annotate(
-                    text, (cx, cy),
-                    fontsize=6, color=color, fontstyle="italic",
+                    text,
+                    (cx, cy),
+                    fontsize=6,
+                    color=color,
+                    fontstyle="italic",
                     bbox=dict(boxstyle="round,pad=0.2", fc="white", ec=color, alpha=0.8),
                 )
             else:
                 ax.annotate(
-                    str(k), (x, y),
-                    textcoords="offset points", xytext=(5, 3),
-                    fontsize=7, color=color,
+                    str(k),
+                    (x, y),
+                    textcoords="offset points",
+                    xytext=(5, 3),
+                    fontsize=7,
+                    color=color,
                 )
 
     ax.set_xlim(0, xlim)

onecomp/lpcd/_gradient_solver.py

@@ -2,6 +2,7 @@
 Copyright 2025-2026 Fujitsu Ltd.
 
 """
+
 from abc import ABC, abstractmethod
 
 import torch
@@ -19,7 +20,7 @@ class PenaltyFn(ABC):
     @abstractmethod
     def __call__(self, module_q: LpcdMetric, module_f: LpcdMetric) -> torch.Tensor:
         pass
- 
+
 
 def gradient_solver(
     lpcd_config: LPCDConfig,
@@ -35,19 +36,21 @@ def gradient_solver(
     device = lpcd_config.device
     batch_size = 16
 
-    assert lpcd_config.gd_batch_size % batch_size == 0, \
-        f"gd_batch_size should be divisible by batch_size, " + \
-        f"but got {lpcd_config.gd_batch_size} and {batch_size}"
-    
-    assert target_modules[0].weight.data.dtype in [torch.float32, torch.bfloat16], \
-        "The model must be loaded in float32 or bfloat16 " + \
-        "for gradient-based LPCD refinement due to numerical stability."
-    
+    assert lpcd_config.gd_batch_size % batch_size == 0, (
+        f"gd_batch_size should be divisible by batch_size, "
+        + f"but got {lpcd_config.gd_batch_size} and {batch_size}"
+    )
+
+    assert target_modules[0].weight.data.dtype in [torch.float32, torch.bfloat16], (
+        "The model must be loaded in float32 or bfloat16 "
+        + "for gradient-based LPCD refinement due to numerical stability."
+    )
+
     # backup and configure grad state
     grad_state_q = [p.requires_grad for p in metric_q.parameters()]
     grad_state_f = [p.requires_grad for p in metric_f.parameters()]
 
-    # set target modules as trainablethe     
+    # set target modules as trainablethe
     metric_q.requires_grad_(False)
     metric_f.requires_grad_(False)
     for module in target_modules:
@@ -60,13 +63,11 @@ def gradient_solver(
 
     optimizer = optim.Adam(metric_q.parameters(), lr=lpcd_config.gd_base_lr)
     scheduler = get_cosine_schedule_with_warmup(
-        optimizer,
-        num_warmup_steps=int(total_iters * 0.1),
-        num_training_steps=total_iters
+        optimizer, num_warmup_steps=int(total_iters * 0.1), num_training_steps=total_iters
     )
 
     # Perform gradient descent to relax the weights
-    for epoch in range(1, lpcd_config.gd_steps+1):
+    for epoch in range(1, lpcd_config.gd_steps + 1):
 
         chunked_inps = zip(
             inps_q.split(batch_size),
@@ -75,7 +76,7 @@ def gradient_solver(
 
         # TODO: shuffle the inputs for better convergence
         for iter, (inp_q, inp_f) in enumerate(chunked_inps):
-        
+
             with torch.no_grad():
                 out_f = metric_f(inp_f.to(device), **kwargs)
 
@@ -95,4 +96,4 @@ def gradient_solver(
     for p, requires_grad in zip(metric_q.parameters(), grad_state_q):
         p.requires_grad_(requires_grad)
     for p, requires_grad in zip(metric_f.parameters(), grad_state_f):
-        p.requires_grad_(requires_grad)
+        p.requires_grad_(requires_grad)

onecomp/lpcd/_lpcd_config.py

@@ -54,4 +54,3 @@ class LPCDConfig:
     gd_batch_size: int = 16
     gd_base_lr: float = 1e-4
     device: str = "cuda:0"
-    

onecomp/lpcd/_lpcd_runner.py

@@ -45,7 +45,7 @@ def run_quantize_with_lpcd(
     lpcd_config: LPCDConfig,
     calibration_config: CalibrationConfig,
 ):
-    """ Run quantization with LPCD.
+    """Run quantization with LPCD.
 
     Args:
         model_config (ModelConfig): Model configuration
@@ -55,8 +55,9 @@ def run_quantize_with_lpcd(
         calibration_config (CalibrationConfig): Calibration configuration
     """
 
-    assert not (qep_config is not None and qep_config.general), \
-        "qep_config.general must be False when qep is enabled."
+    assert not (
+        qep_config is not None and qep_config.general
+    ), "qep_config.general must be False when qep is enabled."
 
     # TODO: Parameterize when necessary
     batch_size = 16
@@ -69,7 +70,7 @@ def run_quantize_with_lpcd(
 
     model_inputs = prepare_calibration_dataset(
         tokenizer=tokenizer,
-        device=torch.device('cpu'),
+        device=torch.device("cpu"),
         calibration_config=calibration_config,
         model=model,
         logger=logger,
@@ -99,7 +100,7 @@ def run_quantize_with_lpcd(
         block_f = copy.deepcopy(block_q)
 
         groups_q = make_grouped_module(block_q, inps_q, kwargs, device)
-        
+
         # Build name→module map for block_f, then align groups_f to
         # groups_q by module name.  Using make_grouped_module on
         # block_f independently can produce a different group ordering
@@ -112,12 +113,17 @@ def run_quantize_with_lpcd(
             name: mod for name, mod in block_q.named_modules() if isinstance(mod, nn.Linear)
         }
         groups_f = [
-            [name_to_module_f[next(n for n, m2 in name_to_module_q.items() if m2 is m)] for m in gq]
+            [
+                name_to_module_f[next(n for n, m2 in name_to_module_q.items() if m2 is m)]
+                for m in gq
+            ]
             for gq in groups_q
         ]
 
         lpcd_metrics = make_lpcd_metrics(lpcd_config, block_q, block_f)
-        lpcd_modules_q = [module for metric, _ in lpcd_metrics.metrics for _, module in metric.named_targets()]
+        lpcd_modules_q = [
+            module for metric, _ in lpcd_metrics.metrics for _, module in metric.named_targets()
+        ]
 
         # 3. For each group of layers, perform the following sequentially
         for group_q, group_f in zip(groups_q, groups_f):
@@ -149,13 +155,10 @@ def run_quantize_with_lpcd(
                 # Skip layers not registered for quantization
                 if module not in quantizer.module_to_name:
                     skipped_name = module_q_to_name.get(module, "<unknown>")
-                    logger.info(
-                        "Skipping layer (not in quantization targets): %s", skipped_name
-                    )
+                    logger.info("Skipping layer (not in quantization targets): %s", skipped_name)
                     continue
                 name = quantizer.module_to_name[module]
 
-
                 # Fall back to standard quantization if the module is not LPCD targets
                 if not module in lpcd_modules_q:
 
@@ -190,10 +193,7 @@ def run_quantize_with_lpcd(
                     # Update the weights of the target layer
                     dtype = module.weight.data.dtype
                     module.weight.data = (
-                        quantizer.results[name]
-                        .compute_dequantized_weight()
-                        .to(device)
-                        .to(dtype)
+                        quantizer.results[name].compute_dequantized_weight().to(device).to(dtype)
                     )
 
                 lpcd_metrics.mark_as_ready(module)
@@ -218,20 +218,11 @@ def run_quantize_with_lpcd(
         inps_f = forward_input(inps_f, block_f, kwargs, batch_size, device)
 
         # DEBUG:Compute MSE between quantized and full-precision outputs
-        mse = compute_mse(
-            block_q,
-            block_f,
-            inps_q,
-            inps_f,
-            batch_size,
-            device,
-            kwargs
-        )
+        mse = compute_mse(block_q, block_f, inps_q, inps_f, batch_size, device, kwargs)
         logger.info(f"[INFO] Layer {block_idx + 1} MSE: {mse:.6e}")
 
         # free memory
         block_q.cpu()
         torch.cuda.empty_cache()
 
     quantizer.execute_post_processing()
-    

onecomp/lpcd/_metric.py

@@ -2,6 +2,7 @@
 Copyright 2025-2026 Fujitsu Ltd.
 
 """
+
 import torch
 from torch import nn
 from transformers.models.llama.modeling_llama import LlamaDecoderLayer
@@ -11,6 +12,8 @@
 from typing import Callable
 
 from ._lpcd_config import LPCDConfig
+
+
 @dataclass
 class ClosedFormSolverArgument:
     lpcd_config: LPCDConfig
@@ -36,12 +39,12 @@ def __init__(self, block: nn.Module):
 
     def is_refineable(self) -> bool:
         return all(self.is_ready.values()) and not self.is_refined
-    
+
     # Implement named targets of the metric
     @abstractmethod
     def named_targets(self) -> list[tuple[str, nn.Module]]:
         pass
-    
+
     # Implement metric computation here
     @abstractmethod
     def forward(self, block_inps: torch.Tensor, kwargs: dict) -> torch.Tensor:
@@ -50,11 +53,10 @@ def forward(self, block_inps: torch.Tensor, kwargs: dict) -> torch.Tensor:
     # (optional) Implement closed-form solvers for specific cases (e.g., llama-out, etc.)
     def closed_form_solvers(self) -> list[Callable[[ClosedFormSolverArgument], None] | None]:
         return [None] * len(self.named_targets())
-    
+
 
 class LpcdMetricGroup:
-    """LpcdMetricGroup organizes LpcdMetrics for a block
-    """
+    """LpcdMetricGroup organizes LpcdMetrics for a block"""
 
     def __init__(self, metrics: list[tuple[LpcdMetric, LpcdMetric]]):
         self.metrics = metrics
@@ -68,31 +70,32 @@ def mark_as_ready(self, module: nn.Module):
             self.module_to_metric[module].is_ready[module] = True
 
     def get_refineable_metrics(self) -> list[LpcdMetric]:
-        return [(metric_q, metric_f) for metric_q, metric_f in self.metrics if metric_q.is_refineable()]
-
+        return [
+            (metric_q, metric_f) for metric_q, metric_f in self.metrics if metric_q.is_refineable()
+        ]
 
 
 from transformers.models.llama.modeling_llama import LlamaDecoderLayer
 from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer
 
+
 def make_lpcd_metrics(
-        lpcd_config: LPCDConfig, 
-        block_q: nn.Module, 
-        block_f: nn.Module
-    ) -> LpcdMetricGroup:
+    lpcd_config: LPCDConfig, block_q: nn.Module, block_f: nn.Module
+) -> LpcdMetricGroup:
     # Implement logic to select and return the appropriate LpcdMetric subclass
     # based on the type of *block* (e.g., LlamaDecoderLayer, etc.)
     if isinstance(block_q, LlamaDecoderLayer):
         # Llama
         from .arch._llama import make_llama_lpcd_metrics
+
         return make_llama_lpcd_metrics(lpcd_config, block_q, block_f)
     elif isinstance(block_q, Qwen3DecoderLayer):
         # Qwen3
         from .arch._qwen3 import make_qwen3_lpcd_metrics
+
         return make_qwen3_lpcd_metrics(lpcd_config, block_q, block_f)
 
     # Add more architecture-specific cases here
     else:
         cls_name = block_q.__class__.__name__
         raise NotImplementedError(f"No LPCD module group implemented for {cls_name}")
-