@@ -279,13 +279,14 @@ def _replace_embedding_with_quantized_group_embedding_for_pre_quantization(
279279 dtype : torch .dtype ,
280280 bit_width : int ,
281281 group_size : Optional [int ] = None ,
282+ scales_precision : torch .dtype = torch .float32 ,
282283):
283284 def filter_fn (child : torch .nn .Module , cur_fqn : str ) -> bool :
284285 # Only replace embedding layers where the checkpoint contains explicit scales
285286 scales_key = f"{ cur_fqn } .scales"
286287 if isinstance (child , nn .Embedding ) and scales_key in checkpoint :
287288 assert checkpoint [f"{ cur_fqn } .weight" ].dtype == torch .int8
288- assert checkpoint [scales_key ].dtype == torch . float32
289+ assert checkpoint [scales_key ].dtype == scales_precision
289290 return True
290291 return False
291292
@@ -297,6 +298,7 @@ def replacement_fn(child: torch.nn.Module) -> torch.nn.Module:
297298 group_size = group_size ,
298299 dtype = dtype ,
299300 packed = False , # TODO(lunwenh): support packed embedding for pre-quantized
301+ scales_precision = scales_precision ,
300302 )
301303 return new_embedding
302304
@@ -309,10 +311,15 @@ def transform_embedding_for_pre_quantization(
309311 dtype : torch .dtype ,
310312 bit_width : int ,
311313 group_size : Optional [int ] = None ,
314+ scales_precision : Optional [torch .dtype ] = None ,
312315) -> torch .nn .Module :
313316 """
314317 Transform the model to be able to load pre-quantized checkpoints that
315318 are quantized with the given bit_width and group size for embedding.
319+
320+ ``scales_precision`` is the dtype the checkpoint stores embedding scales in;
321+ it defaults to ``torch.float32`` to preserve the pre-existing contract
322+ unless explicitly overridden.
316323 """
317324 if group_size is not None and group_size not in [0 , 32 , 64 , 128 , 256 ]:
318325 raise ValueError (
@@ -324,6 +331,7 @@ def transform_embedding_for_pre_quantization(
324331 dtype ,
325332 bit_width ,
326333 group_size ,
334+ scales_precision if scales_precision is not None else torch .float32 ,
327335 )
328336 return module
329337
0 commit comments