MCPcopy Index your code
hub / github.com/NVIDIA/TensorRT-LLM / __init__

Method __init__

tensorrt_llm/quantization/layers.py:248–282  ·  view source on GitHub ↗
(
            self,
            normalized_shape,
            eps=1e-05,
            elementwise_affine=True,
            bias=True,
            dtype=None,
            quant_mode=QuantMode(0),
    )

Source from the content-addressed store, hash-verified

246class SmoothQuantLayerNorm(Module):
247
248 def __init__(
249 self,
250 normalized_shape,
251 eps=1e-05,
252 elementwise_affine=True,
253 bias=True,
254 dtype=None,
255 quant_mode=QuantMode(0),
256 ):
257 super().__init__()
258 if isinstance(normalized_shape, int):
259 normalized_shape = (normalized_shape, )
260 if not quant_mode.has_act_and_weight_quant():
261 raise ValueError(
262 "SmoothQuant layer norm has to have some quantization mode set")
263 self.normalized_shape = tuple(normalized_shape)
264 self.elementwise_affine = elementwise_affine
265 if self.elementwise_affine:
266 self.weight = Parameter(shape=self.normalized_shape, dtype=dtype)
267 if bias:
268 self.bias = Parameter(shape=self.normalized_shape, dtype=dtype)
269 else:
270 self.register_parameter('bias', None)
271 else:
272 self.register_parameter('weight', None)
273 self.register_parameter('bias', None)
274
275 self.eps = eps
276 self.dtype = dtype
277 self.quant_mode = quant_mode
278
279 if self.quant_mode.has_act_and_weight_quant():
280 self.scale_to_int = Parameter(shape=(1, ), dtype=dtype)
281 else:
282 self.register_parameter('scale_to_int', None)
283
284 def forward(self, x):
285 weight = None if self.weight is None else self.weight.value

Callers

nothing calls this directly

Calls 5

QuantModeClass · 0.85
ParameterClass · 0.85
__init__Method · 0.45
register_parameterMethod · 0.45

Tested by

no test coverage detected