pre-commit

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Author: yibinl-nvidia

tensorrt_llm/_torch/models/modeling_starcoder2.py

@@ -23,9 +23,9 @@
 from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.models.modeling_utils import (
-    _load_weights_impl,
     DecoderModel,
     DecoderModelForCausalLM,
+    _load_weights_impl,
     register_auto_model,
 )
 from tensorrt_llm._torch.modules.attention import Attention
@@ -38,7 +38,6 @@
 from tensorrt_llm.functional import PositionEmbeddingType
 
 
-
 class Starcoder2Attention(Attention):
     """
     StarCoder2 Attention with Grouped Query Attention and Sliding Window support.
@@ -122,7 +121,9 @@ def __init__(
                 config=model_config,
             )
         else:
-            raise ValueError(f"Unsupported mlp_type: {config.mlp_type}. Only default (linear) MLP is supported.")
+            raise ValueError(
+                f"Unsupported mlp_type: {config.mlp_type}. Only default (linear) MLP is supported."
+            )
 
         norm_eps = getattr(config, "norm_epsilon", 1e-5)
         self.input_layernorm = LayerNorm(
@@ -219,9 +220,7 @@ def forward(
         lora_params=None,
     ) -> torch.Tensor:
         if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError(
-                "You must specify exactly one of input_ids or inputs_embeds."
-            )
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds.")
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)

tests/integration/defs/accuracy/references/gsm8k.yaml

@@ -275,4 +275,4 @@ bigcode/starcoder2-3b:
 bigcode/starcoder2-7b:
   - accuracy: 26.5
 bigcode/starcoder2-15b:
-  - accuracy: 54.5
+  - accuracy: 54.5

tests/unittest/_torch/modeling/test_modeling_starcoder2.py

@@ -1,8 +1,7 @@
-import pytest
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Any
 
+import pytest
 import torch
 from transformers import Starcoder2Config
 from transformers import Starcoder2ForCausalLM as HFStarcoder2ForCausalLM
@@ -123,6 +122,7 @@ def get_kv_cache_manager(
     )
     return kv_cache_manager
 
+
 @pytest.mark.parametrize(
     "scenario",
     [
@@ -173,13 +173,13 @@ def test_starcoder2_allclose_to_hf(scenario: Scenario) -> None:
         model_config = ModelConfig(pretrained_config=hf_config, attn_backend=backend)
         starcoder2 = Starcoder2ForCausalLM(model_config).to(dtype).to(device).eval()
         starcoder2.load_weights(hf_starcoder2.state_dict())
-    
+
     # Convert LayerNorm random weights to FP32 for numerical stability
     for name, module in starcoder2.named_modules():
         if isinstance(module, LayerNorm):
-            if hasattr(module, 'weight') and module.weight is not None:
+            if hasattr(module, "weight") and module.weight is not None:
                 module.weight.data = module.weight.data.to(torch.float32)
-            if hasattr(module, 'bias') and module.bias is not None:
+            if hasattr(module, "bias") and module.bias is not None:
                 module.bias.data = module.bias.data.to(torch.float32)
 
     num_blocks = 1