[misc] fix model merger (#479)

Author: hiyouga

scripts/model_merger.py

@@ -24,7 +24,7 @@
     AutoConfig,
     AutoModelForCausalLM,
     AutoModelForTokenClassification,
-    AutoModelForVision2Seq,
+    AutoModelForImageTextToText,
     PretrainedConfig,
     PreTrainedModel,
 )
@@ -165,10 +165,10 @@ def process_one_shard(rank, model_state_dict_lst):
 
     if "ForTokenClassification" in architectures[0]:
         AutoClass = AutoModelForTokenClassification
+    elif "ForConditionalGeneration" in architectures[0]:
+        AutoClass = AutoModelForImageTextToText
     elif "ForCausalLM" in architectures[0]:
         AutoClass = AutoModelForCausalLM
-    elif "ForConditionalGeneration" in architectures[0]:
-        AutoClass = AutoModelForVision2Seq
     else:
         raise NotImplementedError(f"Unknown architecture {architectures}.")
 

verl/workers/fsdp_workers.py

@@ -191,14 +191,14 @@ def _build_model_optimizer(
             torch_dtype = PrecisionType.to_dtype(fsdp_config.torch_dtype)
 
         if role == "critic":
-            auto_class = AutoModelForTokenClassification
+            AutoClass = AutoModelForTokenClassification
         elif type(self.model_config) in AutoModelForImageTextToText._model_mapping.keys():
-            auto_class = AutoModelForImageTextToText
+            AutoClass = AutoModelForImageTextToText
         else:
-            auto_class = AutoModelForCausalLM
+            AutoClass = AutoModelForCausalLM
 
         if (not fsdp_config.enable_rank0_init) or self.device_mesh.get_local_rank("fsdp") == 0:
-            model = auto_class.from_pretrained(
+            model = AutoClass.from_pretrained(
                 model_config.model_path,
                 config=self.model_config,
                 torch_dtype=torch_dtype,
@@ -209,7 +209,7 @@ def _build_model_optimizer(
             )
         else:
             with no_init_weights(), init_empty_weights():
-                model = auto_class.from_config(
+                model = AutoClass.from_config(
                     self.model_config,
                     torch_dtype=torch_dtype,
                     attn_implementation="flash_attention_2",