Upload FP8Qwen2ForCausalLM (#6)

- Upload FP8Qwen2ForCausalLM (3e721beb3dfa87636eda6249a856fec8c4dfaae8)


Co-authored-by: Neil <[email protected]>

config.json

@@ -12,7 +12,6 @@
     "AutoModelForTokenClassification": "modeling_fp8_qwen2.FP8Qwen2ForTokenClassification"
   },
   "bos_token_id": 151643,
-  "dtype": "bfloat16",
   "eos_token_id": 151645,
   "fp8_config": {
     "act_block_size": 16,
@@ -68,7 +67,8 @@
   "rope_theta": 1000000.0,
   "sliding_window": null,
   "tie_word_embeddings": false,
-  "transformers_version": "4.57.0",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.54.1",
   "use_cache": true,
   "use_sliding_window": false,
   "vocab_size": 152064

generation_config.json

@@ -2,5 +2,5 @@
   "_from_model_config": true,
   "bos_token_id": 151643,
   "eos_token_id": 151645,
-  "transformers_version": "4.57.0"
+  "transformers_version": "4.54.1"
 }

modeling_fp8_qwen2.py

@@ -47,6 +47,7 @@ from .configuration_fp8_qwen2 import FP8Qwen2Config
 from torchao.float8.float8_training_tensor import Float8TrainingTensor
 
 from quasar.module import (
+    FP8Quant,
     FP8RMSNorm,
     FP8DSLinearWithCoat,
     FP8DSLinearWithCoatWeightBlock,
@@ -65,9 +66,24 @@ class FP8Qwen2MLP(Qwen2MLP):
     def __init__(self, config: FP8Qwen2Config):
         super().__init__(config)
         linear_module = FP8DSLinearWithCoat if config.fp8_config.training_mode else FP8DSLinearWithCoatWeightBlock
-        self.gate_proj = linear_module(self.hidden_size, self.intermediate_size, bias=False, dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"gate_proj"))
-        self.up_proj = linear_module(self.hidden_size, self.intermediate_size, bias=False, dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"up_proj"))
-        self.down_proj = linear_module(self.intermediate_size, self.hidden_size, bias=False, dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"down_proj"))
+        self.gate_proj = linear_module(
+            self.hidden_size,
+            self.intermediate_size,
+            bias=False,
+            dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"gate_proj", scale_dtype=torch.float32),
+        )
+        self.up_proj = linear_module(
+            self.hidden_size,
+            self.intermediate_size,
+            bias=False,
+            dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"up_proj", scale_dtype=torch.float32),
+        )
+        self.down_proj = linear_module(
+            self.intermediate_size,
+            self.hidden_size,
+            bias=False,
+            dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"down_proj", scale_dtype=torch.float32),
+        )
 
         if config.hidden_act == "silu":
             mul_config = FP8MulConfig(
@@ -93,22 +109,46 @@ class FP8Qwen2Attention(Qwen2Attention):
     def __init__(self, config: FP8Qwen2Config, layer_idx: int):
         super().__init__(config, layer_idx)
         linear_module = FP8DSLinearWithCoat if config.fp8_config.training_mode else FP8DSLinearWithCoatWeightBlock
-        self.q_proj = linear_module(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True, dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"q_proj"))
-        self.k_proj = linear_module(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True, dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"k_proj"))
-        self.v_proj = linear_module(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True, dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"v_proj"))
-  
-        if not config.fp8_config.training_mode:
-            # Only when doing inference, we quantize the output of the attention layer.
-            self.o_proj = linear_module(
-                config.num_attention_heads * self.head_dim, config.hidden_size, bias=False,
-                dsgemm_config=FP8DSLinearWithCoatConfig(
-                    fwd_input_quant_type=QuantType.DIV,
-                    layer_name=f"o_proj",
-                    scale_dtype=torch.float32,
-                )
+        self.q_proj = linear_module(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=True,
+            dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"q_proj", scale_dtype=torch.float32),
+        )
+        self.k_proj = linear_module(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=True,
+            dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"k_proj", scale_dtype=torch.float32),
+        )
+        self.v_proj = linear_module(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=True,
+            dsgemm_config=FP8DSLinearWithCoatConfig(layer_name=f"v_proj", scale_dtype=torch.float32),
+        )
+
+        # In both training and inference, we quantize the output of the attention layer.
+        self.o_proj_quant = FP8Quant(
+            quant_config=FP8QuantConfig(
+                float8_dtype=config.fp8_config.float8_dtype,
+                quant_type=QuantType.DIV,
+                fwd_block_size=config.fp8_config.mm_block_size,
+                layer_name=f"o_proj_quant",
+                scale_dtype=torch.float32,
             )
-        
-        
+        )
+        self.o_proj = linear_module(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=False,
+            dsgemm_config=FP8DSLinearWithCoatConfig(
+                fwd_input_quant_type=QuantType.DIV,
+                layer_name=f"o_proj",
+                scale_dtype=torch.float32,
+            ),
+        )
+
     @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
     def forward(
         self,
@@ -157,6 +197,9 @@ class FP8Qwen2Attention(Qwen2Attention):
         )
 
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+
+        # Quantize the output of the attention layer.
+        attn_output = self.o_proj_quant(attn_output)
         attn_output = self.o_proj(attn_output)
         return attn_output, attn_weights
 
@@ -169,8 +212,20 @@ class FP8Qwen2DecoderLayer(GradientCheckpointingLayer):
         self.self_attn = FP8Qwen2Attention(config=config, layer_idx=layer_idx)
 
         self.mlp = FP8Qwen2MLP(config)
-        self.input_layernorm = FP8RMSNorm(config.hidden_size, eps=config.rms_norm_eps, norm_config=FP8RMSNormConfig(mm_block_size=config.fp8_config.mm_block_size, quant_type=QuantType.MUL, save_fp8_input=True))
-        self.post_attention_layernorm = FP8RMSNorm(config.hidden_size, eps=config.rms_norm_eps, norm_config=FP8RMSNormConfig(mm_block_size=config.fp8_config.mm_block_size, quant_type=QuantType.MUL, save_fp8_input=True))
+        self.input_layernorm = FP8RMSNorm(
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+            norm_config=FP8RMSNormConfig(
+                mm_block_size=config.fp8_config.mm_block_size, quant_type=QuantType.MUL, save_fp8_input=True
+            ),
+        )
+        self.post_attention_layernorm = FP8RMSNorm(
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+            norm_config=FP8RMSNormConfig(
+                mm_block_size=config.fp8_config.mm_block_size, quant_type=QuantType.MUL, save_fp8_input=True
+            ),
+        )
         self.attention_type = config.layer_types[layer_idx]
 
     @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
@@ -365,13 +420,13 @@ def make_state_dict_compatible_with_hf(
     """
     # Assert linear keys and undesired linear keys are non-overlapping
     assert set(linear_keys).isdisjoint(set(undesired_linear_keys))
-    
+
     compatible_state_dict = {}
 
     for key in state_dict.keys():
         if any(k in key for k in linear_keys):
             weight = state_dict[key]
-            
+
             if already_fp8:
                 # The name (either weight or weight_scale_inv) is the same as the original key.
                 compatible_state_dict[key] = weight
@@ -381,17 +436,17 @@ def make_state_dict_compatible_with_hf(
                     float8_dtype=config.fp8_config.float8_dtype,
                     quant_type=config.fp8_config.quant_type,
                     fwd_block_size=config.fp8_config.mm_block_size,
-                    scale_dtype=torch.float32, 
+                    scale_dtype=torch.float32,
                 )
                 quant_weight, scale_weight = fp8_quantize_hp2pb(
                     weight, tmp_quant_cfg, block_size=config.fp8_config.mm_block_size
                 )
-                
+
                 name_quant = key.replace("weight", "weight")
                 name_scale = key.replace("weight", "weight_scale_inv")
                 compatible_state_dict[name_quant] = quant_weight
                 compatible_state_dict[name_scale] = scale_weight
-        
+
         elif any(k in key for k in undesired_linear_keys):
             # Dequantize the weight
             if already_fp8:
@@ -400,12 +455,14 @@ def make_state_dict_compatible_with_hf(
                     name_quant = key.replace("weight_scale_inv", "weight")
                     quant_weight = state_dict[name_quant]
                     scale_weight = state_dict[key]
-                    weight = fp8_dequantize_pb2hp(quant_weight, scale_weight, config.fp8_config, block_size=config.fp8_config.mm_block_size)
+                    weight = fp8_dequantize_pb2hp(
+                        quant_weight, scale_weight, config.fp8_config, block_size=config.fp8_config.mm_block_size
+                    )
                     compatible_state_dict[name_quant] = weight
             else:
                 # Do not quantize the weight.
                 compatible_state_dict[key] = state_dict[key]
-                    
+
         else:
             compatible_state_dict[key] = state_dict[key]
     return compatible_state_dict