Fix format.

Signed-off-by: menogrey <1299267905@qq.com>

Author: menogrey

scripts/prepare_hidden_states.py

@@ -507,7 +507,8 @@ def generate(
                 continue
 
             filtered_batch_gpu = {
-                k: v.to(get_local_device(), non_blocking=True) for k, v in filtered_batch.items()
+                k: v.to(get_local_device(), non_blocking=True)
+                for k, v in filtered_batch.items()
             }
             _, _, aux_hidden_states_list, last_hidden_states_list = self.model.extend(
                 **filtered_batch_gpu,

scripts/train_eagle3.py

@@ -73,7 +73,10 @@ def print_cuda_memory_debug(label: str) -> None:
         allocated_bytes = device_module.memory_allocated()
         reserved_bytes = device_module.memory_reserved()
     except Exception as exc:
-        print(f"[memory-debug] {label}: failed to query {device_type} memory: {exc}", flush=True)
+        print(
+            f"[memory-debug] {label}: failed to query {device_type} memory: {exc}",
+            flush=True,
+        )
         return
 
     rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else "NA"
@@ -762,7 +765,10 @@ def run_forward(
             # pixel_values = [pv.to(get_local_device()) for pv in data["pixel_values"]] if args.is_vlm else None
             if args.is_vlm:
                 image_grid_thw = (
-                    [thw.to(get_local_device()).squeeze() for thw in data["image_grid_thw"]]
+                    [
+                        thw.to(get_local_device()).squeeze()
+                        for thw in data["image_grid_thw"]
+                    ]
                     if args.is_vlm
                     else None
                 )
@@ -831,7 +837,9 @@ def run_forward(
             target=target,
             hidden_states=hidden_states,
             position_ids=(
-                data["position_ids"].to(get_local_device()) if "position_ids" in data else None
+                data["position_ids"].to(get_local_device())
+                if "position_ids" in data
+                else None
             ),
             image_grid_thw=image_grid_thw,
             is_vlm=args.is_vlm,

specforge/benchmarks/benchmark_flex_attention.py

@@ -13,7 +13,13 @@
     LlamaFlexAttention,
     prepare_decoder_attention_mask,
 )
-from specforge.utils import empty_cache, get_device_module, get_device_type, get_local_device, synchronize
+from specforge.utils import (
+    empty_cache,
+    get_device_module,
+    get_device_type,
+    get_local_device,
+    synchronize,
+)
 
 dynamo.config.recompile_limit = 64
 

specforge/benchmarks/benchmark_loss.py

@@ -3,8 +3,14 @@
 
 import torch
 
-from specforge.core.loss import log_softmax_loss, _compute_loss
-from specforge.utils import empty_cache, get_device_module, get_device_type, get_local_device, synchronize
+from specforge.core.loss import _compute_loss, log_softmax_loss
+from specforge.utils import (
+    empty_cache,
+    get_device_module,
+    get_device_type,
+    get_local_device,
+    synchronize,
+)
 
 TTT_LENGTH = 7
 
@@ -31,13 +37,20 @@ def benchmark_loss_method(
         target = torch.softmax(
             torch.randn(B, T, V, device=get_local_device(), dtype=torch.float32), dim=-1
         )
-        position_mask = torch.ones((B, T, 1), dtype=torch.bool, device=get_local_device())
+        position_mask = torch.ones(
+            (B, T, 1), dtype=torch.bool, device=get_local_device()
+        )
 
         # Pre-allocate logits tensors for each TTT step
         logits_list = []
         for i in range(TTT_LENGTH):
             logits = torch.randn(
-                B, T, V, device=get_local_device(), requires_grad=True, dtype=torch.float32
+                B,
+                T,
+                V,
+                device=get_local_device(),
+                requires_grad=True,
+                dtype=torch.float32,
             )
             logits_list.append(logits)
 

specforge/core/eagle3.py

@@ -35,7 +35,7 @@
 from specforge.core.lk_loss import compute_acceptance_rate, compute_lk_loss
 from specforge.core.loss import log_softmax_loss
 from specforge.modeling.draft import Eagle3DraftModel
-from specforge.utils import get_compile_backend, empty_cache, padding
+from specforge.utils import empty_cache, get_compile_backend, padding
 
 
 class Eagle3Model(nn.Module):

specforge/distributed.py

@@ -78,7 +78,9 @@ def init_distributed(
         timeout(int): Timeout for collective communication in minutes
         tp_size(int): The degree of tensor parallelism
     """
-    dist.init_process_group(backend=get_dist_backend(), timeout=timedelta(minutes=timeout))
+    dist.init_process_group(
+        backend=get_dist_backend(), timeout=timedelta(minutes=timeout)
+    )
     local_rank = dist.get_rank() % device_count()
     set_device(local_rank)
     print_with_rank(f"bind to device {local_rank}")
@@ -123,7 +125,9 @@ def init_distributed(
     _DP_GROUP = dp_group
     _DRAFT_DP_GROUP = draft_device_mesh.get_group("draft_dp")
     _DRAFT_SP_GROUP = draft_device_mesh.get_group("sp")
-    _DP_DEVICE_MESH = dist.DeviceMesh.from_group(dp_group, device_type=get_device_type())
+    _DP_DEVICE_MESH = dist.DeviceMesh.from_group(
+        dp_group, device_type=get_device_type()
+    )
 
 
 def destroy_distributed():

specforge/modeling/auto.py

@@ -18,6 +18,8 @@
     modeling_utils,
 )
 
+from specforge.utils import get_local_device
+
 from .draft.llama3_eagle import LlamaForCausalLMEagle3
 from .target.custom_backend import (
     GptOssForCausalLM,
@@ -28,7 +30,6 @@
     Qwen3ForCausalLM,
     Qwen3MoeForCausalLM,
 )
-from specforge.utils import get_local_device
 
 
 class AutoEagle3DraftModel(AutoModelForCausalLMBase):

specforge/modeling/draft/flex_attention.py

@@ -77,7 +77,9 @@ def __new__(cls, *args, **kwargs):
     @torch.compiler.disable(recursive=False)
     def __init__(self):
         if not self._is_create_block_mask_compiled:
-            self._compiled_create_block_mask = torch.compile(create_block_mask, backend=get_compile_backend())
+            self._compiled_create_block_mask = torch.compile(
+                create_block_mask, backend=get_compile_backend()
+            )
             self._is_create_block_mask_compiled = True
 
     def __call__(self):