[Misc]Add customized information for models (#4132)

2024-05-01 12:18:14 +08:00
parent c3845d82dc
commit d6f4bd7cdd
9 changed files with 87 additions and 0 deletions
--- a/vllm/model_executor/layers/layernorm.py
+++ b/vllm/model_executor/layers/layernorm.py
@@ -64,3 +64,8 @@ class RMSNorm(nn.Module):
            self.variance_epsilon,
        )
        return out
+
+    def extra_repr(self) -> str:
+        s = f"hidden_size={self.weight.data.size(0)}"
+        s += f", eps={self.variance_epsilon}"
+        return s