ctranslate2 4.6.2__cp314-cp314t-win_amd64.whl → 4.6.3__cp314-cp314t-win_amd64.whl

ctranslate2/converters/fairseq.py

@@ -146,7 +146,9 @@ class FairseqConverter(Converter):
             import_user_module(argparse.Namespace(user_dir=self._user_dir))
 
         with torch.no_grad():
-            checkpoint = checkpoint_utils.load_checkpoint_to_cpu(self._model_path)
+            checkpoint = torch.load(
+                self._model_path, map_location=torch.device("cpu"), weights_only=False
+            )
             args = checkpoint["args"] or checkpoint["cfg"]["model"]
 
             args.data = self._data_dir

ctranslate2/converters/opennmt_py.py

@@ -174,7 +174,9 @@ class OpenNMTPyConverter(Converter):
     def _load(self):
         import torch
 
-        checkpoint = torch.load(self._model_path, map_location="cpu")
+        checkpoint = torch.load(
+            self._model_path, map_location="cpu", weights_only=False
+        )
 
         src_vocabs, tgt_vocabs = get_vocabs(checkpoint["vocab"])
 

ctranslate2/converters/transformers.py

@@ -89,7 +89,7 @@ class TransformersConverter(Converter):
           copy_files: List of filenames to copy from the Hugging Face model to the
             converted model directory.
           load_as_float16: Load the model weights as float16. More precisely, the model
-            will be loaded with ``from_pretrained(..., torch_dtype=torch.float16)``.
+            will be loaded with ``from_pretrained(..., dtype=torch.float16)``.
           revision: Revision of the model to download from the Hugging Face Hub.
           low_cpu_mem_usage: Enable the flag ``low_cpu_mem_usage`` when loading the model
             with ``from_pretrained``.
@@ -123,10 +123,11 @@ class TransformersConverter(Converter):
             tokenizer_class = transformers.AutoTokenizer
 
             kwargs = {
-                "torch_dtype": (
+                "dtype": (
                     torch.float16
                     if self._load_as_float16
-                    else getattr(config, "torch_dtype", None)
+                    else getattr(config, "dtype", None)
+                    or getattr(config, "torch_dtype", None)
                 )
             }
 
@@ -235,7 +236,7 @@ class ModelLoader(abc.ABC):
 
         if isinstance(module, transformers.Conv1D):
             spec.weight = spec.weight.transpose(0, 1)
-        if module.bias is not None:
+        if hasattr(module, "bias") and module.bias is not None:
             spec.bias = module.bias
 
     def set_embeddings(self, spec, module):
@@ -2182,6 +2183,28 @@ class Qwen2Loader(ModelLoader):
             rotary_scaling_type = None
             rotary_scaling_factor = 1
 
+        # Check for AWQ quantization config
+        quantization_config = getattr(model.config, "quantization_config", None)
+        if quantization_config:
+            quant_type = None
+            if quantization_config.quant_method == "awq":
+                quant_type = _SUPPORTED_QUANTIZATION.get(quantization_config.version)
+            if quant_type is None:
+                raise NotImplementedError(
+                    "Quantization type '%s' is not yet implemented. "
+                    "The following Quantization types are currently supported: %s"
+                    % (
+                        quantization_config.quant_method,
+                        ", ".join(_SUPPORTED_QUANTIZATION.keys()),
+                    )
+                )
+            quant_group_size = quantization_config.group_size
+            quant_bits = quantization_config.bits
+        else:
+            quant_type = common_spec.Quantization.CT2
+            quant_group_size = None
+            quant_bits = None
+
         spec = transformer_spec.TransformerDecoderModelSpec.from_config(
             num_layers,
             num_heads,
@@ -2195,9 +2218,12 @@ class Qwen2Loader(ModelLoader):
             rotary_scaling_factor=rotary_scaling_factor,
             rotary_base=getattr(model.config, "rope_theta", 10000),
             num_heads_kv=num_heads_kv,
+            quant_type=quant_type,
+            quant_group_size=quant_group_size,
+            quant_bits=quant_bits,
         )
 
-        self.set_decoder(spec.decoder, model.model)
+        self.set_decoder(spec.decoder, model.model, quant_type)
         self.set_linear(spec.decoder.projection, model.lm_head)
         return spec
 
@@ -2227,7 +2253,7 @@ class Qwen2Loader(ModelLoader):
     def set_layer_norm(self, spec, layer_norm):
         spec.gamma = layer_norm.weight
 
-    def set_decoder(self, spec, module):
+    def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2):
         spec.scale_embeddings = False
         self.set_embeddings(spec.embeddings, module.embed_tokens)
         self.set_layer_norm(spec.layer_norm, module.norm)
@@ -2241,19 +2267,39 @@ class Qwen2Loader(ModelLoader):
             )
 
             split_layers = [common_spec.LinearSpec() for _ in range(3)]
-            self.set_linear(split_layers[0], layer.self_attn.q_proj)
-            self.set_linear(split_layers[1], layer.self_attn.k_proj)
-            self.set_linear(split_layers[2], layer.self_attn.v_proj)
+            self.set_linear(
+                split_layers[0], layer.self_attn.q_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                split_layers[1], layer.self_attn.k_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                split_layers[2], layer.self_attn.v_proj, quant_type=quant_type
+            )
+
+            if quant_type == common_spec.Quantization.CT2:
+                utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers)
+            else:
+                cc_dim = 1 if quant_type == common_spec.Quantization.AWQ_GEMM else 0
+                utils.fuse_linear_prequant(
+                    layer_spec.self_attention.linear[0], split_layers, cc_dim
+                )
 
-            utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers)
             self.set_linear(
                 layer_spec.self_attention.linear[1],
                 layer.self_attn.o_proj,
+                quant_type=quant_type,
             )
 
-            self.set_linear(layer_spec.ffn.linear_0, layer.mlp.gate_proj)
-            self.set_linear(layer_spec.ffn.linear_0_noact, layer.mlp.up_proj)
-            self.set_linear(layer_spec.ffn.linear_1, layer.mlp.down_proj)
+            self.set_linear(
+                layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                layer_spec.ffn.linear_0_noact, layer.mlp.up_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type
+            )
 
             delattr(layer, "self_attn")
             delattr(layer, "mlp")
@@ -2292,6 +2338,28 @@ class Qwen3Loader(ModelLoader):
             rotary_scaling_type = None
             rotary_scaling_factor = 1
 
+        # Check for AWQ quantization config
+        quantization_config = getattr(model.config, "quantization_config", None)
+        if quantization_config:
+            quant_type = None
+            if quantization_config.quant_method == "awq":
+                quant_type = _SUPPORTED_QUANTIZATION.get(quantization_config.version)
+            if quant_type is None:
+                raise NotImplementedError(
+                    "Quantization type '%s' is not yet implemented. "
+                    "The following Quantization types are currently supported: %s"
+                    % (
+                        quantization_config.quant_method,
+                        ", ".join(_SUPPORTED_QUANTIZATION.keys()),
+                    )
+                )
+            quant_group_size = quantization_config.group_size
+            quant_bits = quantization_config.bits
+        else:
+            quant_type = common_spec.Quantization.CT2
+            quant_group_size = None
+            quant_bits = None
+
         spec = transformer_spec.TransformerDecoderModelSpec.from_config(
             num_layers,
             num_heads,
@@ -2307,9 +2375,12 @@ class Qwen3Loader(ModelLoader):
             num_heads_kv=num_heads_kv,
             head_dim=head_dim,
             qk_norm=True,
+            quant_type=quant_type,
+            quant_group_size=quant_group_size,
+            quant_bits=quant_bits,
         )
 
-        self.set_decoder(spec.decoder, model.model)
+        self.set_decoder(spec.decoder, model.model, quant_type)
         self.set_linear(spec.decoder.projection, model.lm_head)
         return spec
 
@@ -2338,7 +2409,7 @@ class Qwen3Loader(ModelLoader):
     def set_layer_norm(self, spec, layer_norm):
         spec.gamma = layer_norm.weight
 
-    def set_decoder(self, spec, module):
+    def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2):
         spec.scale_embeddings = False
         self.set_embeddings(spec.embeddings, module.embed_tokens)
         self.set_layer_norm(spec.layer_norm, module.norm)
@@ -2359,22 +2430,43 @@ class Qwen3Loader(ModelLoader):
             )
 
             split_layers = [common_spec.LinearSpec() for _ in range(3)]
-            self.set_linear(split_layers[0], layer.self_attn.q_proj)
-            self.set_linear(split_layers[1], layer.self_attn.k_proj)
-            self.set_linear(split_layers[2], layer.self_attn.v_proj)
-            utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers)
+            self.set_linear(
+                split_layers[0], layer.self_attn.q_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                split_layers[1], layer.self_attn.k_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                split_layers[2], layer.self_attn.v_proj, quant_type=quant_type
+            )
+
+            if quant_type == common_spec.Quantization.CT2:
+                utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers)
+            else:
+                cc_dim = 1 if quant_type == common_spec.Quantization.AWQ_GEMM else 0
+                utils.fuse_linear_prequant(
+                    layer_spec.self_attention.linear[0], split_layers, cc_dim
+                )
 
             self.set_linear(
                 layer_spec.self_attention.linear[1],
                 layer.self_attn.o_proj,
+                quant_type=quant_type,
             )
 
-            self.set_linear(layer_spec.ffn.linear_0, layer.mlp.gate_proj)
-            self.set_linear(layer_spec.ffn.linear_0_noact, layer.mlp.up_proj)
-            self.set_linear(layer_spec.ffn.linear_1, layer.mlp.down_proj)
+            self.set_linear(
+                layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                layer_spec.ffn.linear_0_noact, layer.mlp.up_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type
+            )
 
             delattr(layer, "self_attn")
             delattr(layer, "mlp")
+            gc.collect()
 
 
 @register_loader("MixFormerSequentialConfig")
@@ -2514,6 +2606,28 @@ class Phi3Loader(ModelLoader):
             rotary_scaling_type = None
             rotary_scaling_factor = 1
 
+        # Check for AWQ quantization config
+        quantization_config = getattr(model.config, "quantization_config", None)
+        if quantization_config:
+            quant_type = None
+            if quantization_config.quant_method == "awq":
+                quant_type = _SUPPORTED_QUANTIZATION.get(quantization_config.version)
+            if quant_type is None:
+                raise NotImplementedError(
+                    "Quantization type '%s' is not yet implemented. "
+                    "The following Quantization types are currently supported: %s"
+                    % (
+                        quantization_config.quant_method,
+                        ", ".join(_SUPPORTED_QUANTIZATION.keys()),
+                    )
+                )
+            quant_group_size = quantization_config.group_size
+            quant_bits = quantization_config.bits
+        else:
+            quant_type = common_spec.Quantization.CT2
+            quant_group_size = None
+            quant_bits = None
+
         spec = transformer_spec.TransformerDecoderModelSpec.from_config(
             num_layers,
             num_heads,
@@ -2529,9 +2643,12 @@ class Phi3Loader(ModelLoader):
             original_max_position_embeddings=original_max_position_embeddings,
             max_position_embeddings=max_position_embeddings,
             num_heads_kv=num_heads_kv,
+            quant_type=quant_type,
+            quant_group_size=quant_group_size,
+            quant_bits=quant_bits,
         )
 
-        self.set_decoder(spec.decoder, model.model)
+        self.set_decoder(spec.decoder, model.model, quant_type)
         self.set_linear(spec.decoder.projection, model.lm_head)
         return spec
 
@@ -2565,7 +2682,7 @@ class Phi3Loader(ModelLoader):
             rotary_scaling_short_factor, dtype=torch.float32
         )
 
-    def set_decoder(self, spec, module):
+    def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2):
         spec.scale_embeddings = False
         self.set_embeddings(spec.embeddings, module.embed_tokens)
         self.set_layer_norm(spec.layer_norm, module.norm)
@@ -2579,9 +2696,15 @@ class Phi3Loader(ModelLoader):
             )
 
             self.set_linear(
-                layer_spec.self_attention.linear[0], layer.self_attn.qkv_proj
+                layer_spec.self_attention.linear[0],
+                layer.self_attn.qkv_proj,
+                quant_type=quant_type,
+            )
+            self.set_linear(
+                layer_spec.self_attention.linear[1],
+                layer.self_attn.o_proj,
+                quant_type=quant_type,
             )
-            self.set_linear(layer_spec.self_attention.linear[1], layer.self_attn.o_proj)
             if (
                 layer.self_attn.rotary_emb.long_factor is not None
                 and layer.self_attn.rotary_emb.short_factor is not None
@@ -2592,10 +2715,30 @@ class Phi3Loader(ModelLoader):
                     layer.self_attn.rotary_emb.short_factor,
                 )
 
-            gate_proj, up_proj = layer.mlp.gate_up_proj.weight.chunk(2, dim=0)
-            layer_spec.ffn.linear_0.weight = gate_proj
-            layer_spec.ffn.linear_0_noact.weight = up_proj
-            self.set_linear(layer_spec.ffn.linear_1, layer.mlp.down_proj)
+            # Handle gate_up_proj differently for AWQ vs regular models
+            if quant_type == common_spec.Quantization.CT2:
+                gate_proj, up_proj = layer.mlp.gate_up_proj.weight.chunk(2, dim=0)
+                layer_spec.ffn.linear_0.weight = gate_proj
+                layer_spec.ffn.linear_0_noact.weight = up_proj
+            else:
+                # AWQ: chunk qweight, scales, and qzeros
+                gate_qweight, up_qweight = layer.mlp.gate_up_proj.qweight.chunk(
+                    2, dim=1
+                )
+                gate_scales, up_scales = layer.mlp.gate_up_proj.scales.chunk(2, dim=1)
+                gate_qzeros, up_qzeros = layer.mlp.gate_up_proj.qzeros.chunk(2, dim=1)
+
+                layer_spec.ffn.linear_0.weight = gate_qweight
+                layer_spec.ffn.linear_0.weight_scale = gate_scales
+                layer_spec.ffn.linear_0.weight_zero = gate_qzeros
+
+                layer_spec.ffn.linear_0_noact.weight = up_qweight
+                layer_spec.ffn.linear_0_noact.weight_scale = up_scales
+                layer_spec.ffn.linear_0_noact.weight_zero = up_qzeros
+
+            self.set_linear(
+                layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type
+            )
 
             delattr(layer, "self_attn")
             delattr(layer, "mlp")
@@ -3325,3 +3468,266 @@ _WHISPER_ALIGNMENT_HEADS = {
         (25, 6),
     ],
 }
+
+
+# Paper: https://arxiv.org/pdf/2504.06225
+@register_loader("T5GemmaConfig")
+class T5GemmaLoader(ModelLoader):
+    @property
+    def architecture_name(self):
+        return "T5GemmaForConditionalGeneration"
+
+    def set_layer_norm(self, spec, layer_norm):
+        spec.gamma = layer_norm.weight.data + 1.0
+
+    def get_model_spec(self, model):
+        encoder_config = model.config.encoder
+        decoder_config = model.config.decoder
+        sliding_window = getattr(model.config, "sliding_window", 4096)
+
+        encoder_num_heads = encoder_config.num_attention_heads
+        encoder_num_heads_kv = getattr(
+            encoder_config, "num_key_value_heads", encoder_num_heads
+        )
+        if encoder_num_heads_kv == encoder_num_heads:
+            encoder_num_heads_kv = None
+
+        encoder = transformer_spec.TransformerEncoderSpec(
+            encoder_config.num_hidden_layers,
+            encoder_config.num_attention_heads,
+            pre_norm=True,
+            activation=_SUPPORTED_ACTIVATIONS[encoder_config.hidden_activation],
+            ffn_glu=True,
+            rms_norm=True,
+            rotary_dim=encoder_config.head_dim,
+            rotary_interleave=False,
+            rotary_base=getattr(encoder_config, "rope_theta", 10000),
+            sliding_window=sliding_window,
+            pre_post_layer_norm=True,
+            num_heads_kv=encoder_num_heads_kv,
+            head_dim=encoder_config.head_dim,
+        )
+
+        decoder_num_heads = decoder_config.num_attention_heads
+        decoder_num_heads_kv = getattr(
+            decoder_config, "num_key_value_heads", decoder_num_heads
+        )
+        if decoder_num_heads_kv == decoder_num_heads:
+            decoder_num_heads_kv = None
+
+        decoder = transformer_spec.TransformerDecoderSpec(
+            decoder_config.num_hidden_layers,
+            decoder_config.num_attention_heads,
+            pre_norm=True,
+            activation=_SUPPORTED_ACTIVATIONS[decoder_config.hidden_activation],
+            ffn_glu=True,
+            rms_norm=True,
+            with_encoder_attention=True,
+            rotary_dim=decoder_config.head_dim,
+            rotary_interleave=False,
+            rotary_base=getattr(decoder_config, "rope_theta", 10000),
+            sliding_window=sliding_window,
+            pre_post_layer_norm=True,
+            external_pre_post_encoder_layers=True,
+            num_heads_kv=decoder_num_heads_kv,
+            head_dim=decoder_config.head_dim,
+        )
+
+        spec = transformer_spec.TransformerSpec(encoder, decoder)
+
+        self.set_encoder(spec.encoder, model.model.encoder, encoder_config)
+
+        self.set_decoder(
+            spec.decoder,
+            model.model.decoder,
+            decoder_config,
+            common_spec.Quantization.CT2,
+        )
+
+        # Tie_word_embeddings
+        self.set_linear(spec.decoder.projection, model.model.decoder.embed_tokens)
+        return spec
+
+    def set_vocabulary(self, spec, tokens):
+        spec.register_source_vocabulary(tokens)
+        spec.register_target_vocabulary(tokens)
+
+    def set_config(self, config, model, tokenizer):
+        config.bos_token = tokenizer.bos_token
+        config.eos_token = tokenizer.eos_token
+        config.unk_token = tokenizer.unk_token
+
+        if hasattr(model.config, "encoder"):
+            config.layer_norm_epsilon = model.config.encoder.rms_norm_eps
+        elif hasattr(model.config, "rms_norm_eps"):
+            config.layer_norm_epsilon = model.config.rms_norm_eps
+        else:
+            config.layer_norm_epsilon = 1e-6
+
+        config.decoder_start_token = tokenizer.bos_token
+
+    def set_encoder(
+        self, spec, encoder, encoder_config, quant_type=common_spec.Quantization.CT2
+    ):
+        spec.scale_embeddings = True
+
+        encoder_emb_spec = (
+            spec.embeddings[0] if isinstance(spec.embeddings, list) else spec.embeddings
+        )
+
+        self.set_embeddings(encoder_emb_spec, encoder.embed_tokens)
+        encoder_emb_spec.multiply_by_sqrt_depth = encoder_config.hidden_size**0.5
+        self.set_layer_norm(spec.layer_norm, encoder.norm)
+
+        module = encoder
+        for i, (layer_spec, layer) in enumerate(zip(spec.layer, module.layers)):
+            self.set_layer_norm(
+                layer_spec.input_layer_norm, layer.pre_self_attn_layernorm
+            )
+            self.set_layer_norm(
+                layer_spec.post_attention_layer_norm, layer.post_self_attn_layernorm
+            )
+
+            # T5GemmaSelfAttention
+            qkv_split_layers = [common_spec.LinearSpec() for _ in range(3)]
+            self.set_linear(
+                qkv_split_layers[0], layer.self_attn.q_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                qkv_split_layers[1], layer.self_attn.k_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                qkv_split_layers[2], layer.self_attn.v_proj, quant_type=quant_type
+            )
+            utils.fuse_linear(layer_spec.self_attention.linear[0], qkv_split_layers)
+            self.set_linear(
+                layer_spec.self_attention.linear[1],
+                layer.self_attn.o_proj,
+                quant_type=quant_type,
+            )
+
+            # T5GemmaRMSNorm
+            self.set_layer_norm(
+                layer_spec.pre_feedforward_layer_norm, layer.pre_feedforward_layernorm
+            )
+            # T5GemmaRMSNorm
+            self.set_layer_norm(
+                layer_spec.post_feedforward_layer_norm, layer.post_feedforward_layernorm
+            )
+
+            # T5GemmaMLP
+            self.set_linear(
+                layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                layer_spec.ffn.linear_0_noact, layer.mlp.up_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type
+            )
+
+            # Clean up
+            delattr(layer, "self_attn")
+            delattr(layer, "mlp")
+            gc.collect()
+
+    def set_decoder(
+        self, spec, module, decoder_config, quant_type=common_spec.Quantization.CT2
+    ):
+        spec.scale_embeddings = True
+        spec.start_from_zero_embedding = False
+
+        self.set_embeddings(spec.embeddings, module.embed_tokens)
+        spec.embeddings.multiply_by_sqrt_depth = decoder_config.hidden_size**0.5
+        self.set_layer_norm(spec.layer_norm, module.norm)
+
+        for i, (layer_spec, layer) in enumerate(zip(spec.layer, module.layers)):
+            # Self-attention block
+            self.set_layer_norm(
+                layer_spec.input_layer_norm, layer.pre_self_attn_layernorm
+            )
+            self.set_layer_norm(
+                layer_spec.post_attention_layer_norm, layer.post_self_attn_layernorm
+            )
+
+            # T5GemmaSelfAttention - QKV projections
+            qkv_split_layers = [common_spec.LinearSpec() for _ in range(3)]
+            self.set_linear(
+                qkv_split_layers[0], layer.self_attn.q_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                qkv_split_layers[1], layer.self_attn.k_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                qkv_split_layers[2], layer.self_attn.v_proj, quant_type=quant_type
+            )
+            utils.fuse_linear(layer_spec.self_attention.linear[0], qkv_split_layers)
+            self.set_linear(
+                layer_spec.self_attention.linear[1],
+                layer.self_attn.o_proj,
+                quant_type=quant_type,
+            )
+
+            # Pre and post cross-attention layer norm
+            self.set_layer_norm(
+                layer_spec.external_pre_encoder_attention_layer_norm,
+                layer.pre_cross_attn_layernorm,
+            )
+
+            self.set_layer_norm(
+                layer_spec.external_post_encoder_attention_layer_norm,
+                layer.post_cross_attn_layernorm,
+            )
+
+            # Cross-attention Q projection
+            self.set_linear(
+                layer_spec.attention.linear[0],
+                layer.cross_attn.q_proj,
+                quant_type=quant_type,
+            )
+
+            # Cross-attention K+V fused
+            kv_split_layers = [common_spec.LinearSpec() for _ in range(2)]
+            self.set_linear(
+                kv_split_layers[0],
+                layer.cross_attn.k_proj,
+                quant_type=quant_type,
+            )
+            self.set_linear(
+                kv_split_layers[1],
+                layer.cross_attn.v_proj,
+                quant_type=quant_type,
+            )
+            utils.fuse_linear(layer_spec.attention.linear[1], kv_split_layers)
+
+            # Cross-attention output projection
+            self.set_linear(
+                layer_spec.attention.linear[2],
+                layer.cross_attn.o_proj,
+                quant_type=quant_type,
+            )
+
+            # Feed-forward block
+            self.set_layer_norm(
+                layer_spec.pre_feedforward_layer_norm, layer.pre_feedforward_layernorm
+            )
+            self.set_layer_norm(
+                layer_spec.post_feedforward_layer_norm, layer.post_feedforward_layernorm
+            )
+
+            # T5GemmaMLP
+            self.set_linear(
+                layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                layer_spec.ffn.linear_0_noact, layer.mlp.up_proj, quant_type=quant_type
+            )
+            self.set_linear(
+                layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type
+            )
+
+            # Clean up
+            delattr(layer, "self_attn")
+            delattr(layer, "cross_attn")
+            delattr(layer, "mlp")
+            gc.collect()

ctranslate2/extensions.py

@@ -556,12 +556,28 @@ def _process_iterable(process_func, iterables, max_batch_size, batch_type, **kwa
 
 def _batch_iterator(iterable, batch_size, batch_type):
     streams = None
-    cur_batch_size = 0
+    max_length = 0
 
     for example in iterable:
         if not isinstance(example, tuple):
             example = (example,)
 
+        if batch_type == "examples":
+            if streams and len(streams[0]) == batch_size:
+                yield streams
+                streams = None
+
+        elif batch_type == "tokens":
+            max_length = max(max_length, len(example[0]))
+
+            if streams and (len(streams[0]) + 1) * max_length > batch_size:
+                yield streams
+                streams = None
+                max_length = len(example[0])
+
+        else:
+            raise ValueError("Invalid batch type %s" % batch_type)
+
         if streams is None:
             streams = tuple([] for _ in example)
         for batch, element in zip(streams, example):
@@ -569,17 +585,5 @@ def _batch_iterator(iterable, batch_size, batch_type):
                 raise ValueError("Input iterables do not have the same length")
             batch.append(element)
 
-        if batch_type == "examples":
-            cur_batch_size += 1
-        elif batch_type == "tokens":
-            cur_batch_size += len(example[0])
-        else:
-            raise ValueError("Invalid batch type %s" % batch_type)
-
-        if cur_batch_size >= batch_size:
-            yield streams
-            streams = None
-            cur_batch_size = 0
-
     if streams is not None:
         yield streams

ctranslate2/specs/attention_spec.py

@@ -34,10 +34,12 @@ class MultiHeadAttentionSpec(model_spec.LayerSpec):
         sliding_window=None,
         qk_norm=False,
         qk_norm_rms=True,
+        has_norm=True,
     ):
         self.queries_scale = model_spec.OPTIONAL
 
-        self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        if has_norm:
+            self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
         self.linear = [
             common_spec.LinearSpec() for _ in range(2 if self_attention else 3)
         ]

ctranslate2/specs/transformer_spec.py

@@ -23,6 +23,16 @@ class TransformerEncoderSpec(model_spec.LayerSpec):
         ffn_glu: bool = False,
         rms_norm: bool = False,
         multi_query_attention: bool = False,
+        num_heads_kv: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        rotary_dim: Optional[int] = None,
+        rotary_interleave: bool = True,
+        rotary_scaling_type: Optional[attention_spec.RotaryScalingType] = None,
+        rotary_scaling_factor: float = 1,
+        rotary_base: float = 10000,
+        sliding_window: Optional[int] = None,
+        qk_norm: Optional[bool] = False,
+        pre_post_layer_norm: bool = False,
     ):
         """Initializes a Transformer encoder specification.
 
@@ -43,8 +53,28 @@ class TransformerEncoderSpec(model_spec.LayerSpec):
           ffn_glu: Use gated linear units in the FFN layers as described in
             https://arxiv.org/abs/2002.05202.
           rms_norm: Use the root mean square layer normalization.
-          multi_query_attention: Use multi-query attention.
+          multi_query_attention: Use multi-query attention (alias for num_heads_kv=1).
+          num_heads_kv: Number of attention heads for the key and value.
+          head_dim: Number of dimensions per attention head.
+          rotary_dim: Apply rotary embeddings to these first N dimensions. If 0, rotary
+            embeddings are applied to all dimensions.
+          rotary_interleave: Interleave the head dimensions when rotary embeddings are applied.
+            Otherwise the head dimensions are sliced in half.
+          rotary_scaling_type: Type of RoPE scaling.
+          rotary_scaling_factor: Factor used in the RoPE scaling.
+          rotary_base: The base period of the rotary embeddings.
+          sliding_window: Max sequence length to retain in KV Cache.
+          qk_norm: Apply layer normalization to the query and key projections.
+          pre_post_layer_norm: Add post layer norm for each pre norm layer.
         """
+
+        if multi_query_attention:
+            if num_heads_kv is not None and num_heads_kv != 1:
+                raise ValueError(
+                    "Enabling multi_query_attention implies num_heads_kv=1"
+                )
+            num_heads_kv = 1
+
         self.multi_query_attention = multi_query_attention
         self.num_heads = np.dtype("int16").type(num_heads)
         self.pre_norm = pre_norm
@@ -60,13 +90,24 @@ class TransformerEncoderSpec(model_spec.LayerSpec):
             self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
         if layernorm_embedding:
             self.layernorm_embedding = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        if sliding_window is not None:
+            self.sliding_window = np.dtype("int32").type(sliding_window)
+
         self.layer = [
             TransformerEncoderLayerSpec(
                 relative_position=relative_position,
                 relative_attention_bias=relative_attention_bias,
                 ffn_glu=ffn_glu,
                 rms_norm=rms_norm,
-                num_heads_kv=1 if multi_query_attention else None,
+                num_heads_kv=num_heads_kv,
+                head_dim=head_dim,
+                rotary_dim=rotary_dim,
+                rotary_interleave=rotary_interleave,
+                rotary_scaling_type=rotary_scaling_type,
+                rotary_scaling_factor=rotary_scaling_factor,
+                rotary_base=rotary_base,
+                qk_norm=qk_norm,
+                pre_post_layer_norm=pre_post_layer_norm,
             )
             for _ in range(num_layers)
         ]
@@ -109,7 +150,8 @@ class TransformerDecoderSpec(model_spec.LayerSpec):
         quant_type: Optional[common_spec.Quantization] = None,
         quant_group_size: Optional[int] = None,
         quant_bits: Optional[int] = None,
-        qk_norm: Optional[bool] = False,
+        qk_norm: bool = False,
+        external_pre_post_encoder_layers: Optional[bool] = False,
     ):
         """Initializes a Transformer decoder specification.
 
@@ -156,6 +198,8 @@ class TransformerDecoderSpec(model_spec.LayerSpec):
           quant_type: quantization type used (like awq... for lower bit quantization)
           quant_group_size: group size of the lower bit quantization
           quant_bits: number of bit of the quantization (ex: 4bit)
+          external_pre_post_encoder_layers: if the encoder attention pre and processing
+            is done outside the attention.
         """
 
         self._config = dict()
@@ -172,12 +216,6 @@ class TransformerDecoderSpec(model_spec.LayerSpec):
                 )
             num_heads_kv = 1
 
-        if with_encoder_attention and num_heads_kv not in (None, 1, num_heads):
-            raise ValueError(
-                "num_heads_kv=%d is not supported in the cross-attention layers"
-                % num_heads_kv
-            )
-
         self.num_heads = np.dtype("int16").type(num_heads)
         self.pre_norm = pre_norm
         self.activation = np.dtype("int8").type(activation)
@@ -224,6 +262,7 @@ class TransformerDecoderSpec(model_spec.LayerSpec):
                 head_dim=head_dim,
                 sliding_window=sliding_window,
                 qk_norm=qk_norm,
+                external_pre_post_encoder_layers=external_pre_post_encoder_layers,
             )
             for _ in range(num_layers)
         ]
@@ -254,7 +293,15 @@ class TransformerEncoderLayerSpec(model_spec.LayerSpec):
         ffn_glu=False,
         rms_norm=False,
         num_heads_kv=None,
+        head_dim=None,
         sliding_window=None,
+        rotary_dim: Optional[int] = None,
+        rotary_interleave: bool = True,
+        rotary_scaling_type: Optional[attention_spec.RotaryScalingType] = None,
+        rotary_scaling_factor: float = 1,
+        rotary_base: float = 10000,
+        qk_norm=False,
+        pre_post_layer_norm: bool = False,
     ):
         self.self_attention = attention_spec.MultiHeadAttentionSpec(
             self_attention=True,
@@ -262,10 +309,32 @@ class TransformerEncoderLayerSpec(model_spec.LayerSpec):
             relative_attention_bias=relative_attention_bias,
             rms_norm=rms_norm,
             num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
             sliding_window=sliding_window,
+            rotary_dim=rotary_dim,
+            rotary_interleave=rotary_interleave,
+            rotary_scaling_type=rotary_scaling_type,
+            rotary_scaling_factor=rotary_scaling_factor,
+            rotary_base=rotary_base,
+            qk_norm=qk_norm,
         )
         self.ffn = FeedForwardSpec(glu=ffn_glu, rms_norm=rms_norm)
 
+        if pre_post_layer_norm:
+            self.input_layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
+            self.post_attention_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+            self.pre_feedforward_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+            self.post_feedforward_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+
+            delattr(self.self_attention, "layer_norm")
+            delattr(self.ffn, "layer_norm")
+
 
 class TransformerDecoderLayerSpec(model_spec.LayerSpec):
     def __init__(
@@ -289,6 +358,7 @@ class TransformerDecoderLayerSpec(model_spec.LayerSpec):
         head_dim=None,
         sliding_window=None,
         qk_norm=False,
+        external_pre_post_encoder_layers=False,
     ):
         self.self_attention = attention_spec.MultiHeadAttentionSpec(
             self_attention=True,
@@ -312,8 +382,10 @@ class TransformerDecoderLayerSpec(model_spec.LayerSpec):
             self.attention = attention_spec.MultiHeadAttentionSpec(
                 rms_norm=rms_norm,
                 num_heads_kv=num_heads_kv,
+                head_dim=head_dim,
                 sliding_window=sliding_window,
                 qk_norm=qk_norm,
+                has_norm=external_pre_post_encoder_layers is False,
             )
 
         self.ffn = FeedForwardSpec(glu=ffn_glu, rms_norm=rms_norm)
@@ -329,10 +401,21 @@ class TransformerDecoderLayerSpec(model_spec.LayerSpec):
             delattr(self.ffn, "layer_norm")
 
         if pre_post_layer_norm:
+            # Self-attention layer norms
             self.input_layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
             self.post_attention_layer_norm = common_spec.LayerNormSpec(
                 rms_norm=rms_norm
             )
+
+            if with_encoder_attention and external_pre_post_encoder_layers:
+                self.external_post_encoder_attention_layer_norm = (
+                    common_spec.LayerNormSpec(rms_norm=rms_norm)
+                )
+                self.external_pre_encoder_attention_layer_norm = (
+                    common_spec.LayerNormSpec(rms_norm=rms_norm)
+                )
+
+            # Feed-forward layer norms
             self.pre_feedforward_layer_norm = common_spec.LayerNormSpec(
                 rms_norm=rms_norm
             )
@@ -562,7 +645,7 @@ class TransformerDecoderModelSpec(model_spec.LanguageModelSpec):
         quant_type: Optional[common_spec.Quantization] = None,
         quant_group_size: Optional[int] = None,
         quant_bits: Optional[int] = None,
-        qk_norm: Optional[bool] = False,
+        qk_norm: bool = False,
     ):
         """Creates a Transformer decoder model specification.
 

ctranslate2/version.py

@@ -1,3 +1,3 @@
 """Version information."""
 
-__version__ = "4.6.2"
+__version__ = "4.6.3"

{ctranslate2-4.6.2.dist-info → ctranslate2-4.6.3.dist-info}/METADATA

@@ -1,9 +1,10 @@
 Metadata-Version: 2.4
 Name: ctranslate2
-Version: 4.6.2
+Version: 4.6.3
 Summary: Fast inference engine for Transformer models
 Home-page: https://opennmt.net
 Author: OpenNMT
+License: MIT
 Project-URL: Documentation, https://opennmt.net/CTranslate2
 Project-URL: Forum, https://forum.opennmt.net
 Project-URL: Gitter, https://gitter.im/OpenNMT/CTranslate2
@@ -13,7 +14,6 @@ Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: GPU :: NVIDIA CUDA :: 12 :: 12.4
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: 3.9
@@ -34,6 +34,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: keywords
+Dynamic: license
 Dynamic: project-url
 Dynamic: requires-dist
 Dynamic: requires-python
@@ -49,7 +50,7 @@ The project implements a custom runtime that applies many performance optimizati
 
 The following model types are currently supported:
 
-* Encoder-decoder models: Transformer base/big, M2M-100, NLLB, BART, mBART, Pegasus, T5, Whisper
+* Encoder-decoder models: Transformer base/big, M2M-100, NLLB, BART, mBART, Pegasus, T5, Whisper T5Gemma
 * Decoder-only models: GPT-2, GPT-J, GPT-NeoX, OPT, BLOOM, MPT, Llama, Mistral, Gemma, CodeGen, GPTBigCode, Falcon, Qwen2
 * Encoder-only models: BERT, DistilBERT, XLM-RoBERTa
 
@@ -160,6 +161,16 @@ Executed with 4 threads on a [*c5.2xlarge*](https://aws.amazon.com/ec2/instance-
 
 Executed with CUDA 11 on a [*g5.xlarge*](https://aws.amazon.com/ec2/instance-types/g5/) Amazon EC2 instance equipped with a NVIDIA A10G GPU (driver version: 510.47.03).
 
+## Contributing
+
+CTranslate2 is a community-driven project. We welcome contributions of all kinds:
+* **New Model Support:** Help us implement more Transformer architectures.
+* **Performance:** Propose optimizations for CPU or GPU kernels.
+* **Bug Reports:** Open an issue if you find something not working as expected.
+* **Documentation:** Improve our guides or add new examples.
+
+Check out our [Contributing Guide](CONTRIBUTING.md) to learn how to set up your development environment.
+
 ## Additional resources
 
 * [Documentation](https://opennmt.net/CTranslate2)

{ctranslate2-4.6.2.dist-info → ctranslate2-4.6.3.dist-info}/RECORD

@@ -1,33 +1,33 @@
 ctranslate2/__init__.py,sha256=CGqShDaFxQ-u-aCtVq99T4HKuBdMB8b49l2KSxnQb8M,1735
-ctranslate2/_ext.cp314t-win_amd64.pyd,sha256=ryUxQaSEW8-dMZYWQsHYQXoi-2dmaTbqkSnFIeC3z0I,774656
-ctranslate2/ctranslate2.dll,sha256=9zIz4dY3yV1kTTKaipyQwjcGDwzZ3OzKiOkNpXdcQ1U,58389504
-ctranslate2/cudnn64_9.dll,sha256=wHzEfy-kpWZZPHr0qn5X7fCamFoP3dFMuNb0VuJSrwU,438840
-ctranslate2/extensions.py,sha256=axO2FI8ddiFmlko2AzQ6VcdtF-3hDA7VmPGnTIkrPkI,21782
+ctranslate2/_ext.cp314t-win_amd64.pyd,sha256=SJzyI2umi8fNXsvbshO65rCof-2B7QCfPf0t44_NEGY,775168
+ctranslate2/ctranslate2.dll,sha256=9TRGHvoyNSzXe9eEc3jKJa5-21-AeSENgp3DXvhCZ4M,58590720
+ctranslate2/cudnn64_9.dll,sha256=ntvN_3OwrwcOsWCyzmbln-ygSqAXNR2O7cxejhSZZ9I,266288
+ctranslate2/extensions.py,sha256=kDNt0H9KvfNCc3PrRGzfkj9Fkvna84i2O5Y-rav6UkU,21940
 ctranslate2/libiomp5md.dll,sha256=mCIzNmsK_NoeD1WgsTQJfjW3eWE_VN22nmhebNBrdV8,1614192
 ctranslate2/logging.py,sha256=P9evHdxuMx_iHvwJjEASEq-j5062H64Pl5-fJjxEuHk,1221
-ctranslate2/version.py,sha256=f2Hk9NHTYgXftujV8JVkeOzenykZ9QzbsZ-nIt9U1uc,53
+ctranslate2/version.py,sha256=TboXlbA67GNmSOm1v2u_U8AKgYh5iminMMLTvi3Xho4,53
 ctranslate2/converters/__init__.py,sha256=ufYjcXf2sK4fiXAUU6tIJyWmNuLjKFf_KH3GWLXe4ls,507
 ctranslate2/converters/converter.py,sha256=Qkb8NGLLmgqMT6HZkFq61zwbxyq3NlWcaxLZ6Ap-YOQ,3601
 ctranslate2/converters/eole_ct2.py,sha256=RUcDJH_2AUt0jDs5oAqccE6tQPbO9LQ6JmVriC1DTy8,12564
-ctranslate2/converters/fairseq.py,sha256=uQpd-ftYSO4c6WdEwCUyuZWhzWX1UTG7dGOC6EtcDVE,12765
+ctranslate2/converters/fairseq.py,sha256=2vlBk4AVCHwXxKkwPHVmcjyfo1dAV0_DJS1i6q-44NE,12822
 ctranslate2/converters/marian.py,sha256=1_7P3EbIDPOdyJbtb_Lp-LCBPBb9A8E9OhzoyFwTb64,11274
 ctranslate2/converters/openai_gpt2.py,sha256=1rXKM2ZURZHWRv4XZ135fPkVWpM4rTG-q7VR7OD6d-A,3304
-ctranslate2/converters/opennmt_py.py,sha256=Vva60az6tGqlQXs0UgC09r_fCD3u2u6wUJB-8V4OUFQ,13183
+ctranslate2/converters/opennmt_py.py,sha256=zex4TbHiiJMy0tkqQg39oNjxmSZKf8dnRLH3iQ1H4z0,13227
 ctranslate2/converters/opennmt_tf.py,sha256=uBRp2wz5xriSQcA_c0S0ekY7ws6RpRX_0EKeMRdM7-s,16222
 ctranslate2/converters/opus_mt.py,sha256=5KbPaTiBhhorPzMpTugIfIJ8SgcqHfJUbJrWKBN-Djs,1254
-ctranslate2/converters/transformers.py,sha256=zwqUFFFwLpam6z5lpBz2rgfYj065CbsdT9S_xVqPjCk,126110
+ctranslate2/converters/transformers.py,sha256=VRal3vKSQrAOvcNPwewjVMtgvWskz0KD5bdIrpNrZNA,142380
 ctranslate2/converters/utils.py,sha256=w7NG39lx-9dOdL57OqKVTdC__opkuP8RACg1TLlUJwM,3817
 ctranslate2/models/__init__.py,sha256=53p98uemtuvVPz8xK7_LbOhBiUJJu-c-NdmOHJgdXus,497
 ctranslate2/specs/__init__.py,sha256=9GabtSyczznYqiqUS6XvULi8pQ3_3RNRogXobGP0G80,653
-ctranslate2/specs/attention_spec.py,sha256=0JhCBrbb20G07UFnUAYIUtfcqn4VtflJHYWGIunwKDw,3442
+ctranslate2/specs/attention_spec.py,sha256=FnaSiQREWQw_cURgsCb9_aIpGOCxyVGTCpIOdd-08v8,3492
 ctranslate2/specs/common_spec.py,sha256=freTDhQMy5PYofBrij4_FDgrKokMYApWSPIpASZIlJc,1608
 ctranslate2/specs/model_spec.py,sha256=atCAYzDEIzyJ1TCayFGZVutHqSWa1ww-vbZ0OiIJqh8,25736
-ctranslate2/specs/transformer_spec.py,sha256=43jOIvCSbAvqZJ1IyvRdGUa4f-zhdKhQBOXvp0T8YLE,30360
+ctranslate2/specs/transformer_spec.py,sha256=-GJ0oSjI3ns-Ei_-xXIM_P2GaZxt5Z-g03zJ0m_4ciU,34317
 ctranslate2/specs/wav2vec2_spec.py,sha256=NITsuOuf2F5bU1-aXit8-WEtWV9fH2Eq7A7857UyYho,2106
 ctranslate2/specs/wav2vec2bert_spec.py,sha256=UgtsJWC9mMgJ7bn4T_xg1uXK0rqA4-9tT2KMGVgPKnw,3529
 ctranslate2/specs/whisper_spec.py,sha256=_vm1sc5yOowOJ4iyvcxMXrgt-UcLJrZT8OtPscUXcQQ,2447
-ctranslate2-4.6.2.dist-info/METADATA,sha256=r5HnmZE0BMI60j3N0GmDdM6l7Q7KW3w5nLLOX_AKCRY,10354
-ctranslate2-4.6.2.dist-info/WHEEL,sha256=IxxXYqBIlKEzFy9ulBJ928Gdqg6XQ3DHti4avqq3myk,102
-ctranslate2-4.6.2.dist-info/entry_points.txt,sha256=ZHkojut_TmVRHl0bJIGm2b9wqr98GAJqxN9rlJtQshs,466
-ctranslate2-4.6.2.dist-info/top_level.txt,sha256=1hUaWzcFIuSo2BAIUHFA3Osgsu6S1giq0y6Rosv8HOQ,12
-ctranslate2-4.6.2.dist-info/RECORD,,
+ctranslate2-4.6.3.dist-info/METADATA,sha256=awoc6t4JSxpv51lmfAG28ZG91FhGQ8DHspyLzLqLo_Q,10839
+ctranslate2-4.6.3.dist-info/WHEEL,sha256=IxxXYqBIlKEzFy9ulBJ928Gdqg6XQ3DHti4avqq3myk,102
+ctranslate2-4.6.3.dist-info/entry_points.txt,sha256=ZHkojut_TmVRHl0bJIGm2b9wqr98GAJqxN9rlJtQshs,466
+ctranslate2-4.6.3.dist-info/top_level.txt,sha256=1hUaWzcFIuSo2BAIUHFA3Osgsu6S1giq0y6Rosv8HOQ,12
+ctranslate2-4.6.3.dist-info/RECORD,,