[misc] Update qwen_moe flops counter (#522)

Co-authored-by: Yaowei Zheng <hiyouga@buaa.edu.cn>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: 3 people

verl/models/monkey_patch.py

@@ -12,12 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 
 from ..utils.py_functional import is_transformers_version_greater_than
 from .transformers.flash_attention_utils import flash_attention_forward
-from .transformers.qwen2_vl import qwen2_vl_base_forward, qwen2_vl_model_forward
 
 
 SUPPORTED_MODEL_TYPE = (
@@ -55,6 +53,8 @@ def apply_ulysses_patch(model_type: str) -> None:
         )
         from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLForConditionalGeneration, Qwen2VLModel
 
+        from .transformers.qwen2_vl import qwen2_vl_base_forward, qwen2_vl_model_forward
+
         # fix text-image mixed data
         Qwen2VLModel.forward = qwen2_vl_base_forward
         Qwen2_5_VLModel.forward = qwen2_vl_base_forward
@@ -63,8 +63,16 @@ def apply_ulysses_patch(model_type: str) -> None:
         Qwen2_5_VLForConditionalGeneration.forward = qwen2_vl_model_forward
     elif model_type in QWEN3_VL_MODELS:
         from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration, Qwen3VLModel
+        from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import (
+            Qwen3VLMoeForConditionalGeneration,
+            Qwen3VLMoeModel,
+        )
+
+        from .transformers.qwen3_vl import qwen3_vl_base_forward, qwen3_vl_model_forward
 
         # fix text-image mixed data
-        Qwen3VLModel.forward = qwen2_vl_base_forward
+        Qwen3VLModel.forward = qwen3_vl_base_forward
+        Qwen3VLMoeModel.forward = qwen3_vl_base_forward
         # TODO: add linear cross entropy kernels
-        Qwen3VLForConditionalGeneration.forward = qwen2_vl_model_forward
+        Qwen3VLForConditionalGeneration.forward = qwen3_vl_model_forward
+        Qwen3VLMoeForConditionalGeneration.forward = qwen3_vl_model_forward

verl/utils/flops_counter.py

@@ -66,9 +66,13 @@ def __init__(self, config: "LlamaConfig"):
         _ESTIMATE_FUNC = {
             "llama": self._estimate_llama_flops,
             "qwen2": self._estimate_llama_flops,
+            "qwen2_moe": self._estimate_qwen2_moe_flops,
             "qwen2_vl": self._estimate_llama_flops,
             "qwen2_5_vl": self._estimate_llama_flops,
             "qwen3": self._estimate_llama_flops,
+            "qwen3_vl": self._estimate_llama_flops,
+            "qwen3_moe": self._estimate_qwen2_moe_flops,
+            "qwen3_vl_moe": self._estimate_qwen2_moe_flops,
         }
 
         if config.model_type not in _ESTIMATE_FUNC:
@@ -115,6 +119,44 @@ def _estimate_llama_flops(self, tokens_sum: int, batch_seqlens: List[int], delta
         flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12
         return flops_achieved
 
+    def _estimate_qwen2_moe_flops(self, tokens_sum: int, batch_seqlens: List[int], delta_time: float) -> float:
+        config = self.config.text_config if hasattr(self.config, "text_config") else self.config
+        hidden_size = config.hidden_size
+        vocab_size = config.vocab_size
+        num_hidden_layers = config.num_hidden_layers
+        num_key_value_heads = config.num_key_value_heads
+        num_attention_heads = config.num_attention_heads
+        moe_intermediate_size = config.moe_intermediate_size
+        moe_topk = config.num_experts_per_tok
+        num_experts = config.num_experts
+
+        head_dim = getattr(config, "head_dim", hidden_size // num_attention_heads)
+        q_size = num_attention_heads * head_dim
+        k_size = num_key_value_heads * head_dim
+        v_size = num_key_value_heads * head_dim
+
+        # non-attn per layer parm
+        # gate + moe export
+        moe_mlp_N = hidden_size * moe_topk * moe_intermediate_size * 3 + hidden_size * num_experts
+        attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim)
+        emd_and_lm_head_N = vocab_size * hidden_size * 2
+        # non-attn all_layer parm
+        dense_N = (moe_mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N
+        # non-attn all_layer & all_token fwd & bwd flops
+        dense_N_flops = 6 * dense_N * tokens_sum
+
+        # attn all_layer & all_token fwd & bwd flops
+        seqlen_square_sum = 0
+        for seqlen in batch_seqlens:
+            seqlen_square_sum += seqlen * seqlen
+
+        attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers
+
+        # all_layer & all_token fwd & bwd flops
+        flops_all_token = dense_N_flops + attn_qkv_flops
+        flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12
+        return flops_achieved
+
     def estimate_flops(self, batch_seqlens: List[int], delta_time: float) -> Tuple[float, float]:
         """
         Estimate the FLOPS based on the number of valid tokens in the current batch and the time taken.

verl/utils/logger/logger.py

@@ -69,6 +69,12 @@ def __init__(self, config: dict[str, Any]) -> None:
         with open(os.path.join(config["trainer"]["save_checkpoint_path"], "experiment_config.json"), "w") as f:
             json.dump(config, f, indent=2)
 
+        with open(os.path.join(config["trainer"]["save_checkpoint_path"], "experiment_log.jsonl"), "w") as f:
+            pass
+
+        with open(os.path.join(config["trainer"]["save_checkpoint_path"], "generations.log"), "w") as f:
+            pass
+
     def log(self, data: dict[str, Any], step: int) -> None:
         with open(os.path.join(self.config["trainer"]["save_checkpoint_path"], "experiment_log.jsonl"), "a") as f:
             f.write(json.dumps({"step": step, **unflatten_dict(data)}) + "\n")