style: apply pre-commit hooks (isort, black) to global_ptq

Ran `pre-commit install` and `pre-commit run --all-files` as requested
in review.  isort and black (line-length 99) reformat the `global_ptq/`
package, which had not been run through the hooks before.  No
functional changes -- formatting only.

Because the hooks run over every file, this also touches modules
outside the scope of this PR: `_core/dbf_adapter.py`,
`_core/gptq_adapter.py`, `_core/losses.py`,
`example/example_global_ptq.py` and the two test modules under
`global_ptq/tests/`.

The other hooks (no-japanese, copyright-header, no-email-address) pass
without changes.

Author: fujisawa-yoshihiko

global_ptq/example/example_global_ptq.py

@@ -5,9 +5,10 @@
 
 """
 
-from onecomp import Runner, ModelConfig, GPTQ, CalibrationConfig, setup_logger
 from onecomp_globalptq import GlobalPTQ
 
+from onecomp import GPTQ, CalibrationConfig, ModelConfig, Runner, setup_logger
+
 MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 setup_logger()

global_ptq/onecomp_globalptq/global_ptq/_core/core.py

@@ -20,13 +20,31 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .dbf_adapter import (
+    load_dbf_state,
+    restore_dbf_original,
+    save_dbf_state,
+    setup_dbf_differentiable,
+    setup_dbf_forwards_only,
+    write_back_dbf_binary,
+    write_back_dbf_scaling,
+)
+from .gptq_adapter import (
+    find_gptq_modules,
+    load_gptq_state,
+    restore_gptq_original,
+    save_gptq_state,
+    setup_gptq_differentiable,
+    setup_gptq_forwards_only,
+    write_back_gptq_params,
+)
 from .helpers import (
     detect_quantization_method,
     disable_gradient_checkpointing,
     enable_gradient_checkpointing,
-    remove_input_require_grads,
     get_language_model_backbone,
     get_logits,
+    remove_input_require_grads,
 )
 from .losses import (
     clear_hooks,
@@ -36,24 +54,6 @@
     remove_hooks,
     setup_intermediate_hooks,
 )
-from .gptq_adapter import (
-    find_gptq_modules,
-    load_gptq_state,
-    restore_gptq_original,
-    save_gptq_state,
-    setup_gptq_differentiable,
-    setup_gptq_forwards_only,
-    write_back_gptq_params,
-)
-from .dbf_adapter import (
-    load_dbf_state,
-    restore_dbf_original,
-    save_dbf_state,
-    setup_dbf_differentiable,
-    setup_dbf_forwards_only,
-    write_back_dbf_binary,
-    write_back_dbf_scaling,
-)
 from .mdbf_adapter import (
     load_mdbf_state,
     restore_mdbf_original,
@@ -272,8 +272,7 @@ def build_param_to_layer_map(
         for p in mod.parameters():
             if p.requires_grad:
                 param_to_layer[id(p)] = layer_idx
-        for attr in ("_opt_scales", "_opt_zeros", "_opt_intweight",
-                     "_opt_bp1", "_opt_bp3"):
+        for attr in ("_opt_scales", "_opt_zeros", "_opt_intweight", "_opt_bp1", "_opt_bp3"):
             if hasattr(mod, attr):
                 param_to_layer[id(getattr(mod, attr))] = layer_idx
 
@@ -333,8 +332,7 @@ def compute_fisher_diagonal(
             pid = id(p)
             if pid in param_to_layer:
                 fisher_acc[param_to_layer[pid]] = (
-                    fisher_acc.get(param_to_layer[pid], 0.0)
-                    + p.grad.float().pow(2).sum().item()
+                    fisher_acc.get(param_to_layer[pid], 0.0) + p.grad.float().pow(2).sum().item()
                 )
 
         model.zero_grad()
@@ -398,8 +396,7 @@ def set_layer_grad(
         if lidx in layer_idx_set:
             for p in mod.parameters():
                 p.requires_grad_(enable)
-            for attr in ("_opt_scales", "_opt_zeros", "_opt_intweight",
-                         "_opt_bp1", "_opt_bp3"):
+            for attr in ("_opt_scales", "_opt_zeros", "_opt_intweight", "_opt_bp1", "_opt_bp3"):
                 if hasattr(mod, attr):
                     getattr(mod, attr).requires_grad_(enable)
 
@@ -464,7 +461,10 @@ def eval_kl(
         logits_s = get_logits(model(input_ids))
         logits_t = _teacher_logits(teacher_model, input_ids, teacher_dev, dev)
         total += compute_kl_loss(
-            logits_t, logits_s, temperature, attention_mask=attention_mask,
+            logits_t,
+            logits_s,
+            temperature,
+            attention_mask=attention_mask,
         ).item()
         n += 1
     if was_training:
@@ -649,9 +649,7 @@ def run_kl_distillation(
 
     The model is modified **in-place**.  Returns a results dict.
     """
-    dev = torch.device(
-        student_device or ("cuda" if torch.cuda.is_available() else "cpu")
-    )
+    dev = torch.device(student_device or ("cuda" if torch.cuda.is_available() else "cpu"))
     teacher_dev = torch.device(teacher_device) if teacher_device else dev
 
     # ------------------------------------------------------------------
@@ -682,8 +680,11 @@ def run_kl_distillation(
     # ------------------------------------------------------------------
     logger.info("Loading calibration data (n=%d, len=%d)...", num_calibration_samples, max_length)
     dataloader = _prepare_dataloader(
-        model_config, num_calibration_samples, max_length,
-        calibration_strategy, calibration_seed,
+        model_config,
+        num_calibration_samples,
+        max_length,
+        calibration_strategy,
+        calibration_seed,
         calibration_dataset=calibration_dataset,
         model=quantized_model,
     )
@@ -715,7 +716,10 @@ def run_kl_distillation(
     if method == "gptq":
         gptq_modules = detected_modules
         original_forwards, scaling_params, intweight_params = setup_gptq_differentiable(
-            gptq_modules, dev, gptq_optimize_intweight, ste_k,
+            gptq_modules,
+            dev,
+            gptq_optimize_intweight,
+            ste_k,
         )
         param_groups = [{"params": scaling_params, "lr": gptq_lr}]
         if intweight_params:
@@ -730,7 +734,8 @@ def run_kl_distillation(
     elif method == "dbf":
         dbf_modules = detected_modules
         original_forwards, scaling_params, binary_params = setup_dbf_differentiable(
-            dbf_modules, optimize_binary,
+            dbf_modules,
+            optimize_binary,
         )
         all_dbf_params = list(scaling_params)
         if binary_params:
@@ -746,7 +751,9 @@ def run_kl_distillation(
     elif method == "mdbf":
         mdbf_modules = detected_modules
         original_forwards, scaling_params, binary_params = setup_mdbf_differentiable(
-            mdbf_modules, optimize_binary, ste_k=mdbf_ste_k,
+            mdbf_modules,
+            optimize_binary,
+            ste_k=mdbf_ste_k,
         )
         logger.info("MDBF binary STE sharpness mdbf_ste_k=%.4g", mdbf_ste_k)
         all_mdbf_params = list(scaling_params)
@@ -783,27 +790,35 @@ def run_kl_distillation(
         logger.info("Computing Fisher diagonal for per-layer LR...")
         ptl_map = build_param_to_layer_map(detected_modules)
         fisher_diag = compute_fisher_diagonal(
-            quantized_model, dataloader, dev, ptl_map, n_samples=fisher_n_samples,
+            quantized_model,
+            dataloader,
+            dev,
+            ptl_map,
+            n_samples=fisher_n_samples,
         )
         fisher_mult = build_fisher_lr_multipliers(fisher_diag, fisher_min_mult, fisher_max_mult)
         if fisher_mult:
             logger.info(
                 "Fisher LR multipliers: min=%.3f, max=%.3f, layers=%d",
-                min(fisher_mult.values()), max(fisher_mult.values()), len(fisher_mult),
+                min(fisher_mult.values()),
+                max(fisher_mult.values()),
+                len(fisher_mult),
             )
             # Determine default weight decay based on method (AdamW vs Adam)
             default_wd = 0.01 if method == "gptq" else 0.0
-            
+
             new_groups: list = []
             for pg in param_groups:
                 for p in pg["params"]:
                     pid = id(p)
                     mult = fisher_mult.get(ptl_map.get(pid, -1), 1.0)
-                    new_groups.append({
-                        "params": [p],
-                        "lr": pg["lr"] * mult,
-                        "weight_decay": pg.get("weight_decay", default_wd),
-                    })
+                    new_groups.append(
+                        {
+                            "params": [p],
+                            "lr": pg["lr"] * mult,
+                            "weight_decay": pg.get("weight_decay", default_wd),
+                        }
+                    )
             param_groups = new_groups
 
     # ------------------------------------------------------------------
@@ -842,7 +857,8 @@ def run_kl_distillation(
             set_layer_grad(detected_modules, all_layer_indices - initial_active, False)
             logger.info(
                 "Progressive unfreeze: %d/%d layers active initially",
-                len(initial_active), total_layers,
+                len(initial_active),
+                total_layers,
             )
 
     # ------------------------------------------------------------------
@@ -884,7 +900,9 @@ def run_kl_distillation(
     if use_lr_schedule:
         scheduler = torch.optim.lr_scheduler.LambdaLR(
             optimizer,
-            lr_lambda=lambda step: cosine_warmup_lr_lambda(step, effective_total_steps, warmup_steps, min_lr_ratio),
+            lr_lambda=lambda step: cosine_warmup_lr_lambda(
+                step, effective_total_steps, warmup_steps, min_lr_ratio
+            ),
         )
 
     # EMA tracker
@@ -932,7 +950,12 @@ def run_kl_distillation(
         restore_mdbf_original(mdbf_modules, original_forwards)
 
     initial_kl = eval_kl(
-        quantized_model, teacher_model, dataloader, dev, temperature, teacher_dev,
+        quantized_model,
+        teacher_model,
+        dataloader,
+        dev,
+        temperature,
+        teacher_dev,
     )
     logger.info("Initial KL = %.6f", initial_kl)
 
@@ -974,9 +997,8 @@ def run_kl_distillation(
 
             n_accum_steps += 1
             is_accum_boundary = (
-                (batch_idx + 1) % grad_accum_steps == 0
-                or batch_idx == total_batches - 1
-            )
+                batch_idx + 1
+            ) % grad_accum_steps == 0 or batch_idx == total_batches - 1
             # Actual number of steps in this accumulation cycle
             current_accum_steps = n_accum_steps if is_accum_boundary else grad_accum_steps
 
@@ -989,11 +1011,17 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
                 with amp_ctx:
                     logits_s = get_logits(quantized_model(input_ids))
                     logits_t = _teacher_logits(
-                        teacher_model, input_ids, teacher_dev, dev,
+                        teacher_model,
+                        input_ids,
+                        teacher_dev,
+                        dev,
                     )
 
                 kl = compute_kl_loss(
-                    logits_t, logits_s, temperature, attention_mask=attention_mask,
+                    logits_t,
+                    logits_s,
+                    temperature,
+                    attention_mask=attention_mask,
                 )
 
                 inter = torch.tensor(0.0, device=dev)
@@ -1057,7 +1085,7 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
 
             # --- Standard path ---
             else:
-                if (batch_idx % grad_accum_steps == 0):
+                if batch_idx % grad_accum_steps == 0:
                     optimizer.zero_grad()
 
                 loss, kl_loss, inter_loss = _forward_and_loss()
@@ -1085,7 +1113,7 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
 
                     if scheduler is not None:
                         scheduler.step()
-                    
+
                     optimizer.zero_grad()
                     n_accum_steps = 0
 
@@ -1094,7 +1122,9 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
                 ema_tracker.update(all_opt_params)
 
             epoch_kl += kl_loss.item()
-            epoch_inter += inter_loss.item() if isinstance(inter_loss, torch.Tensor) else inter_loss
+            epoch_inter += (
+                inter_loss.item() if isinstance(inter_loss, torch.Tensor) else inter_loss
+            )
             epoch_loss += loss.item()
             n_batches += 1
 
@@ -1119,7 +1149,12 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
                 restore_mdbf_original(mdbf_modules, original_forwards)
 
             current_kl = eval_kl(
-                quantized_model, teacher_model, dataloader, dev, temperature, teacher_dev,
+                quantized_model,
+                teacher_model,
+                dataloader,
+                dev,
+                temperature,
+                teacher_dev,
             )
 
             if current_kl < best_kl:
@@ -1221,7 +1256,12 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     # Final evaluation
     quantized_model.eval()
     final_kl = eval_kl(
-        quantized_model, teacher_model, dataloader, dev, temperature, teacher_dev,
+        quantized_model,
+        teacher_model,
+        dataloader,
+        dev,
+        temperature,
+        teacher_dev,
     )
 
     # Cleanup
@@ -1232,7 +1272,7 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         restore_dbf_original(dbf_modules, original_forwards, cleanup=True)
     elif method == "mdbf":
         restore_mdbf_original(mdbf_modules, original_forwards, cleanup=True)
-    
+
     del teacher_model
     gc.collect()
     torch.cuda.empty_cache()