Fix: Skip batch size check for pixel_values in pipelining · Pull Request !1233 · Ascend/MindSpeed-MM

1. Fix: Skip batch size check for pixel_values in pipelining
---
## Motivation

When running inference with models like GLM-4.1V-9B-Thinking, which handle high-resolution images, under a pipeline parallel setup, the program crashes. The crash occurs in parallel_wrapper.py due to a batch size consistency check that fails for the pixel_values tensor.

Error Log:

```bash
Exception: All values in model_forward_kwargs must have the same first dimension (batch size). Expected batch size: 1, but tensor 'pixel_values' has shape torch.Size([14308, 1176]).
All tensor shapes: {'input_ids': torch.Size([1, 3596]), ..., 'pixel_values': torch.Size([14308, 1176]), ...}
```

## Modification

The root cause is that for models like GLM-4.1V-9B-Thinking/Qwen2-VL, the pixel_values argument passed to the core model is not the raw image tensor but the patch embeddings generated by the Vision Encoder. Its shape is [num_patches, hidden_size], so its first dimension is not the batch_size.

The modification updates the batch size check logic in the _with_pipelining_forward_step method within mindspeed_mm/tasks/inference/pipeline/parallel_wrapper.py. Specifically, pixel_values is added to the list of keys to be excluded from the batch size consistency check. This aligns its handling with other special-cased tensors like position_ids and cache_position.

**Code Change:**
```diff
- first_dims = [v.shape[0] for k, v in model_forward_kwargs.items() if (k != "position_ids" and k != "cache_position" and v is not None)]
+ first_dims = [v.shape[0] for k, v in model_forward_kwargs.items() if (k != "position_ids" and k != "cache_position" and k != "pixel_values" and v is not None)]
```

## Self-test (Optional)

After applying this modification, I successfully ran inference with the GLM-4.1V-9B-Thinking/Qwen2-VL model using the same configuration that previously caused the crash. The batch size exception is no longer raised, and the model generates output correctly.

## BC-breaking (Optional)

This change is not backward-compatibility breaking. It only relaxes a check for a specific keyword argument and does not alter any APIs or existing model behaviors.

## Checklist

**Before PR**:

- The new code needs to comply with the Clean Code specification.
 - The PR content is self-checked, and the expression can be clear and the writing standardized.

**After PR**:

- [ ] CLA has been signed and all committers have signed the CLA in this PR.
- [ ] The ci-pipeline is passed, Code Check is passed.

类型	指派人员	状态
审查	王姜奔	已完成 (0/0人)
测试	王姜奔	已完成 (0/0人)

Ascend/MindSpeed-MM暂停 .gitee-modal { width: 500px !important; }

搜索帮助

Ascend/MindSpeed-MM
暂停