.

Author: ZiyiTsang

verl/models/transformers/glm4v.py

@@ -42,24 +42,6 @@
 logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
 
 
-if is_flash_attn_2_available():
-    from flash_attn import flash_attn_func, flash_attn_varlen_func
-
-    _flash_supports_window_size = "window_size" in inspect.signature(flash_attn_func).parameters
-    _flash_supports_deterministic = "deterministic" in inspect.signature(flash_attn_func).parameters
-    _flash_use_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
-# if is_npu_available:
-#     from transformers.integrations.npu_flash_attention import npu_flash_attn_func as flash_attn_func
-#     from transformers.integrations.npu_flash_attention import npu_flash_attn_varlen_func as flash_attn_varlen_func
-#     from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask
-
-#     _flash_supports_window_size = "window_size" in inspect.signature(flash_attn_func).parameters
-#     _flash_supports_deterministic = "deterministic" in inspect.signature(flash_attn_func).parameters
-#     _flash_use_top_left_mask = flash_attn_supports_top_left_mask()
-
-_flash_deterministic_enabled = os.getenv("FLASH_ATTENTION_DETERMINISTIC", "0") == "1"
-
 
 def get_rope_index(
     processor,
@@ -181,353 +163,3 @@ def get_rope_index(
     return position_ids
 
 
-def prepare_fa2_from_position_ids(
-    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, position_ids: torch.Tensor
-):
-    assert position_ids.ndim == 2  # (batch_size, seq_length)
-    query = query.contiguous().view(-1, query.size(-2), query.size(-1))
-    key = key.contiguous().view(-1, key.size(-2), key.size(-1))
-    value = value.contiguous().view(-1, value.size(-2), value.size(-1))
-    position_ids = position_ids.view(-1)
-    cu_seqlens = torch.cat(
-        (
-            (position_ids == 0).nonzero().view(-1).to(torch.int32),
-            torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32),
-        )
-    )
-    max_length = cu_seqlens.diff().max()  # use cu_seqlens to infer max_length for qwen2vl mrope
-    return (query, key, value, (cu_seqlens, cu_seqlens), (max_length, max_length))
-
-
-def _custom_flash_attention_forward(
-    query_states: torch.Tensor,
-    key_states: torch.Tensor,
-    value_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor],
-    query_length: int,
-    is_causal: bool = True,
-    position_ids: Optional[torch.Tensor] = None,
-    use_top_left_mask: bool = False,
-    deterministic: Optional[bool] = None,
-    **kwargs,
-):
-    """
-    Patches flash attention forward to handle 3D position ids in mrope. (3, batch_size, seq_length)
-    """
-    # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length).
-    flash_kwargs = {}
-
-    if _flash_supports_deterministic:
-        flash_kwargs["deterministic"] = deterministic if deterministic is not None else _flash_deterministic_enabled
-
-    if kwargs.get("softcap") is not None:
-        flash_kwargs["softcap"] = kwargs.pop("softcap")
-
-    query_states, key_states, value_states = fa_peft_integration_check(
-        query_states, key_states, value_states, target_dtype=torch.bfloat16
-    )
-
-    if position_ids is not None:
-        assert position_ids.ndim == 2  # (batch_size, seq_length / sp_size)
-
-    sp_size = get_ulysses_sequence_parallel_world_size()
-    if sp_size > 1:
-        # qkv: (batch_size, seq_length / sp_size, num_head, head_size)
-        validate_ulysses_config(query_states.size(2), sp_size)
-        query_states = gather_seq_scatter_heads(query_states, seq_dim=1, head_dim=2)
-        key_states = gather_seq_scatter_heads(key_states, seq_dim=1, head_dim=2)
-        value_states = gather_seq_scatter_heads(value_states, seq_dim=1, head_dim=2)
-        position_ids_lst = [torch.empty_like(position_ids) for _ in range(sp_size)]
-        position_ids = dist.all_gather(position_ids_lst, position_ids, group=get_ulysses_sequence_parallel_group())
-        position_ids = torch.cat(position_ids_lst, dim=-1)  # (batch_size, seq_length)
-
-    if position_ids is not None and query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all():
-        batch_size = query_states.size(0)
-        q, k, v, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = prepare_fa2_from_position_ids(
-            query_states, key_states, value_states, position_ids
-        )
-        attn_output = flash_attn_varlen_func(
-            q=q,
-            k=k,
-            v=v,
-            cu_seqlens_q=cu_seqlens_q,
-            cu_seqlens_k=cu_seqlens_k,
-            max_seqlen_q=max_seqlen_q,
-            max_seqlen_k=max_seqlen_k,
-            dropout_p=kwargs.pop("dropout", 0.0),
-            softmax_scale=kwargs.pop("softmax_scale", None),
-            causal=is_causal,
-            **flash_kwargs,
-        )
-        attn_output = attn_output.view(batch_size, -1, attn_output.size(-2), attn_output.size(-1))
-    else:
-        attn_output = _flash_attention_forward(
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            query_length,
-            is_causal=is_causal,
-            use_top_left_mask=use_top_left_mask,
-            deterministic=deterministic,
-            **kwargs,
-        )  # do not pass position_ids to old flash_attention_forward
-
-    if sp_size > 1:
-        # (batch_size, seq_length, num_head, head_size)
-        attn_output = gather_heads_scatter_seq(attn_output, head_dim=2, seq_dim=1)
-
-    return attn_output
-
-
-def glm4v_attn_forward(
-    self: "Glm4vTextAttention",
-    hidden_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
-    **kwargs,
-) -> tuple[torch.Tensor, None, None]:
-    from transformers.models.glm4v.modeling_glm4v import apply_multimodal_rotary_pos_emb, repeat_kv
-
-    bsz, q_len, _ = hidden_states.size()  # q_len = seq_length / sp_size
-    query_states = self.q_proj(hidden_states)  # (batch_size, seq_length / sp_size, num_heads * head_size)
-    key_states = self.k_proj(hidden_states)
-    value_states = self.v_proj(hidden_states)
-
-    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
-    # Because the input can be padded, the absolute sequence length depends on the max position id.
-    cos, sin = position_embeddings
-    query_states, key_states = apply_multimodal_rotary_pos_emb(
-        query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]
-    )
-    key_states = repeat_kv(key_states, self.num_key_value_groups)
-    value_states = repeat_kv(value_states, self.num_key_value_groups)
-    dropout_rate = 0.0 if not self.training else self.attention_dropout
-
-    # This is before the transpose
-    q_len = query_states.shape[2]
-
-    # FA2 uses non-transposed inputs
-    query_states = query_states.transpose(1, 2)
-    key_states = key_states.transpose(1, 2)
-    value_states = value_states.transpose(1, 2)
-
-    attn_output = _custom_flash_attention_forward(
-        query_states,
-        key_states,
-        value_states,
-        attention_mask,
-        query_length=q_len,
-        is_causal=getattr(self, "is_causal", True),
-        dropout=dropout_rate,
-        use_top_left_mask=_flash_use_top_left_mask,
-        position_ids=position_ids,  # important: pass position ids
-    )  # (batch_size, seq_length / sp_size, num_head, head_size)
-    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
-    attn_output = self.o_proj(attn_output)
-    return attn_output, None
-
-
-def _get_input_embeds(
-    model: "Glm4vForConditionalGeneration",
-    input_ids: torch.LongTensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    pixel_values: Optional[torch.FloatTensor] = None,
-    pixel_values_videos: Optional[torch.FloatTensor] = None,
-    image_grid_thw: Optional[torch.LongTensor] = None,
-    video_grid_thw: Optional[torch.LongTensor] = None,
-):
-    inputs_embeds = model.get_input_embeddings()(input_ids)
-    if pixel_values is not None:
-        pixel_values = pixel_values.type(model.visual.dtype)
-        image_embeds = model.visual(pixel_values, grid_thw=image_grid_thw)
-        n_image_tokens = (input_ids == model.config.image_token_id).sum().item()
-        n_image_features = image_embeds.shape[0]
-        if n_image_tokens != n_image_features:
-            raise ValueError(
-                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
-            )
-
-        mask = input_ids == model.config.image_token_id
-        mask_unsqueezed = mask.unsqueeze(-1)
-        mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
-        image_mask = mask_expanded.to(inputs_embeds.device)
-
-        image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
-        inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
-
-    if pixel_values_videos is not None:
-        pixel_values_videos = pixel_values_videos.type(model.visual.dtype)
-        video_embeds = model.visual(pixel_values_videos, grid_thw=video_grid_thw)
-        n_video_tokens = (input_ids == model.config.video_token_id).sum().item()
-        n_video_features = video_embeds.shape[0]
-        if n_video_tokens != n_video_features:
-            raise ValueError(
-                f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
-            )
-
-        mask = input_ids == model.config.video_token_id
-        mask_unsqueezed = mask.unsqueeze(-1)
-        mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
-        video_mask = mask_expanded.to(inputs_embeds.device)
-
-        video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
-        inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
-
-    if model.training and pixel_values is None and pixel_values_videos is None:  # handle mixed text-image data
-        pixel_values = torch.zeros((16, 1176), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
-        image_grid_thw = torch.tensor([[1, 4, 4]], dtype=torch.long, device=inputs_embeds.device)
-        image_embeds = model.visual(pixel_values, grid_thw=image_grid_thw)
-        inputs_embeds += 0.0 * image_embeds.mean()
-
-    if attention_mask is not None:
-        attention_mask = attention_mask.to(inputs_embeds.device)
-
-    return inputs_embeds, attention_mask
-
-
-def process_position_ids(position_ids: torch.Tensor) -> torch.Tensor:
-    if position_ids.ndim != 3 or position_ids.size(0) != 4:
-        # we concat the text position ids with the 3D vision position ids by default
-        # see https://github.com/huggingface/transformers/pull/39447
-        raise ValueError("position_ids should be a 3D tensor of shape (4, batch_size, seq_length).")
-
-    return position_ids
-
-
-@dataclass
-class Glm4vCausalLMOutputForPPO(Glm4vCausalLMOutputWithPast):
-    log_probs: Optional[torch.FloatTensor] = None
-    entropy: Optional[torch.FloatTensor] = None
-
-
-def glm4v_base_forward(
-    self: "Glm4vForConditionalGeneration",
-    input_ids: torch.LongTensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    labels: Optional[torch.LongTensor] = None,
-    pixel_values: Optional[torch.FloatTensor] = None,
-    pixel_values_videos: Optional[torch.FloatTensor] = None,
-    image_grid_thw: Optional[torch.LongTensor] = None,
-    video_grid_thw: Optional[torch.LongTensor] = None,
-    **kwargs,
-):
-    kwargs["inputs_embeds"], kwargs["attention_mask"] = _get_input_embeds(
-        self, input_ids, attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw
-    )  # avoid lora module having multiple keyword arguments
-    return self.language_model(
-        input_ids=None,
-        **kwargs,
-    )
-
-
-def glm4v_forward(
-    self: "Glm4vForConditionalGeneration",
-    input_ids: torch.LongTensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    pixel_values: Optional[torch.FloatTensor] = None,
-    pixel_values_videos: Optional[torch.FloatTensor] = None,
-    image_grid_thw: Optional[torch.LongTensor] = None,
-    video_grid_thw: Optional[torch.LongTensor] = None,
-    **kwargs,
-):
-    return self.model(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        position_ids=process_position_ids(position_ids),
-        pixel_values=pixel_values,
-        pixel_values_videos=pixel_values_videos,
-        image_grid_thw=image_grid_thw,
-        video_grid_thw=video_grid_thw,
-        **kwargs,
-    )
-
-
-def forward_with_normal_backend(
-    self: Glm4vForConditionalGeneration,
-    input_ids: torch.LongTensor = None,
-    labels: Optional[torch.LongTensor] = None,
-    temperature: float = 1.0,
-    **kwargs,
-) -> "Glm4vCausalLMOutputWithPast":
-    outputs = glm4v_forward(self, input_ids, **kwargs)
-    hidden_states = outputs[0]
-    logits = self.lm_head(hidden_states)
-
-    return Glm4vCausalLMOutputWithPast(
-        logits=logits,
-        hidden_states=outputs.hidden_states,
-    )
-
-
-def forward_with_torch_backend(
-    self: Glm4vForConditionalGeneration,
-    input_ids: torch.LongTensor = None,
-    labels: Optional[torch.LongTensor] = None,
-    temperature: float = 1.0,
-    **kwargs,
-) -> tuple | Glm4vCausalLMOutputForPPO:
-    from verl.utils.experimental.torch_functional import FusedLinearForPPO
-
-    outputs = glm4v_forward(self, input_ids, **kwargs)
-    hidden_states = outputs[0]
-
-    # Loss calculations
-    if labels is not None:
-        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)
-    elif input_ids is not None:
-        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)
-    else:
-        raise RuntimeError("To use forward_with_torch_backend, either labels or input_ids must be provided.")
-
-    fused_linear_for_ppo = FusedLinearForPPO()
-    log_probs, entropy = fused_linear_for_ppo.forward(
-        hidden_states=hidden_states,
-        vocab_weights=self.lm_head.weight,
-        input_ids=rolled_labels,
-        temperature=temperature,
-    )
-    return Glm4vCausalLMOutputForPPO(
-        log_probs=log_probs,
-        entropy=entropy,
-        hidden_states=outputs.hidden_states,
-    )
-
-
-def forward_with_triton_backend(
-    self: Glm4vForConditionalGeneration,
-    input_ids: torch.LongTensor = None,
-    labels: Optional[torch.LongTensor] = None,
-    temperature: float = 1.0,
-    **kwargs,
-) -> tuple | Glm4vCausalLMOutputForPPO:
-    from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy
-
-    outputs = glm4v_forward(self, input_ids, **kwargs)
-    hidden_states = outputs[0]
-
-    # Loss calculations
-    if labels is not None:
-        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)
-    elif input_ids is not None:
-        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)
-    else:
-        raise RuntimeError("To use forward_with_triton_backend, either labels or input_ids must be provided.")
-
-    log_probs, entropy = linear_cross_entropy(
-        hidden_states,
-        self.lm_head.weight,
-        rolled_labels,
-        temperature,
-        "none",
-    )
-    return Glm4vCausalLMOutputForPPO(
-        log_probs=log_probs,
-        entropy=entropy,
-        hidden_states=outputs.hidden_states,
-    )